🎉 Initial commit

mqtt-client/logging_setup.py

@@ -0,0 +1,142 @@
+"""Logging setup for the MQTT client.
+
+Supports:
+  - Console logging (always).
+  - File logging (optional, controlled by LOG_TO_FILE).
+  - Loki push (optional, controlled by LOKI_URL). Loki pushes are batched
+    and sent from a background thread so a slow or unavailable Loki never
+    blocks ingestion.
+"""
+
+import logging
+import os
+import queue
+import sys
+import threading
+from logging.handlers import RotatingFileHandler
+from typing import Optional
+
+import requests
+from pythonjsonlogger import jsonlogger
+
+LOKI_URL = os.getenv("LOKI_URL", "").strip()
+LOKI_USER = os.getenv("LOKI_USER", "")
+LOKI_PASSWORD = os.getenv("LOKI_PASSWORD", "")
+LOKI_FLUSH_SECONDS = float(os.getenv("LOKI_FLUSH_SECONDS", "5"))
+LOKI_TIMEOUT_SECONDS = float(os.getenv("LOKI_TIMEOUT_SECONDS", "5"))
+
+
+class LokiHandler(logging.Handler):
+    """Buffer log records and push them to Loki in batches.
+
+    ``emit()`` only enqueues. A daemon thread flushes the buffer every
+    ``flush_interval`` seconds, grouping records by level into one stream
+    per level so a single HTTP request carries many entries.
+    """
+
+    _MAX_BATCH = 500
+
+    def __init__(
+        self,
+        url: str,
+        user: str = "",
+        password: str = "",
+        flush_interval: Optional[float] = None,
+        timeout: Optional[float] = None,
+    ) -> None:
+        super().__init__()
+        self.url = url.rstrip("/") + "/loki/api/v1/push"
+        self.auth = (user, password) if user else None
+        self._flush_interval = (
+            LOKI_FLUSH_SECONDS if flush_interval is None else flush_interval
+        )
+        self._timeout = LOKI_TIMEOUT_SECONDS if timeout is None else timeout
+        self._buffer: queue.SimpleQueue[logging.LogRecord] = queue.SimpleQueue()
+        self._stop = threading.Event()
+        self._thread = threading.Thread(
+            target=self._publish_loop, name="loki-publisher", daemon=True
+        )
+        self._thread.start()
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self._buffer.put(record)
+
+    def _publish_loop(self) -> None:
+        while not self._stop.wait(self._flush_interval):
+            self.flush()
+
+    def flush(self) -> None:
+        records: list[logging.LogRecord] = []
+        while len(records) < self._MAX_BATCH:
+            try:
+                records.append(self._buffer.get_nowait())
+            except queue.Empty:
+                break
+        if not records:
+            return
+
+        streams: dict[str, list[list[str]]] = {}
+        for record in records:
+            level = record.levelname.lower()
+            streams.setdefault(level, []).append(
+                [str(int(record.created * 1_000_000_000)), self.format(record)]
+            )
+        # Loki requires increasing timestamps within a stream.
+        for values in streams.values():
+            values.sort(key=lambda value: int(value[0]))
+
+        payload = {
+            "streams": [
+                {
+                    "stream": {"service": "mqtt-client", "level": level},
+                    "values": values,
+                }
+                for level, values in streams.items()
+            ]
+        }
+        try:
+            response = requests.post(
+                self.url, json=payload, auth=self.auth, timeout=self._timeout
+            )
+            response.raise_for_status()
+        except Exception:  # noqa: BLE001 - never let logging break ingestion
+            pass
+
+    def close(self) -> None:
+        self._stop.set()
+        self.flush()
+        super().close()
+
+
+def setup_logging() -> logging.Logger:
+    level = os.getenv("LOG_LEVEL", "INFO").upper()
+    logger = logging.getLogger("mqtt-client")
+    logger.setLevel(level)
+    logger.handlers.clear()
+
+    formatter = jsonlogger.JsonFormatter(
+        "%(asctime)s %(levelname)s %(name)s %(message)s"
+    )
+
+    console = logging.StreamHandler(sys.stdout)
+    console.setFormatter(formatter)
+    logger.addHandler(console)
+
+    if os.getenv("LOG_TO_FILE", "true").lower() == "true":
+        path = os.getenv("LOG_FILE_PATH", "/var/log/mqtt-client/mqtt-client.log")
+        try:
+            os.makedirs(os.path.dirname(path), exist_ok=True)
+            file_handler = RotatingFileHandler(
+                path, maxBytes=10 * 1024 * 1024, backupCount=5
+            )
+            file_handler.setFormatter(formatter)
+            logger.addHandler(file_handler)
+        except OSError as exc:
+            # Logging must never prevent startup (e.g. root-owned or
+            # read-only log directory).
+            logger.warning("File logging disabled: %s", exc)
+
+    if LOKI_URL:
+        logger.addHandler(LokiHandler(LOKI_URL, LOKI_USER, LOKI_PASSWORD))
+
+    return logger