🎉 Initial commit

mqtt-client/.dockerignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.pyc
+*.log
+.env
+.venv/

mqtt-client/Dockerfile

@@ -0,0 +1,17 @@
+# Python 3.13+ is required: ssl.SSLContext.set_psk_client_callback (TLS-PSK)
+# is only available from 3.13.
+FROM python:3.13-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+RUN mkdir -p /var/log/mqtt-client \
+    && useradd --system --no-create-home mqtt \
+    && chown -R mqtt:mqtt /var/log/mqtt-client
+USER mqtt
+
+CMD ["python", "-u", "main.py"]

mqtt-client/db_writer.py

@@ -0,0 +1,185 @@
+"""Database writer for persisting MQTT telemetry into TimescaleDB.
+
+Messages are enqueued by the MQTT callback thread and flushed to the
+database in batches by a dedicated worker thread. This keeps blocking I/O
+off the MQTT network loop and amortizes commit overhead over many rows.
+"""
+
+import logging
+import os
+import queue
+import threading
+import time
+from datetime import datetime, timezone
+from typing import Any, Optional
+
+import psycopg
+from psycopg.types.json import Jsonb
+
+logger = logging.getLogger("mqtt-client")
+
+_INSERT_SQL = """
+    INSERT INTO telemetry (time, topic, device_id, payload)
+    VALUES (%s, %s, %s, %s)
+"""
+
+
+class DbWriter:
+    """Queue-backed batch writer for telemetry rows.
+
+    ``write()`` only enqueues, so it is safe to call from the MQTT network
+    loop. A worker thread groups rows into batches and inserts each batch
+    in a single transaction; failed batches are retried and re-enqueued
+    rather than dropped.
+    """
+
+    def __init__(
+        self,
+        batch_size: Optional[int] = None,
+        flush_interval: Optional[float] = None,
+        queue_limit: int = 10_000,
+    ) -> None:
+        if batch_size is None:
+            batch_size = int(os.getenv("DB_BATCH_SIZE", "500"))
+        if flush_interval is None:
+            flush_interval = float(os.getenv("DB_FLUSH_SECONDS", "1.0"))
+        self._batch_size = max(1, batch_size)
+        self._flush_interval = max(0.05, flush_interval)
+        self._queue: queue.Queue[tuple[datetime, str, Any]] = queue.Queue(
+            maxsize=queue_limit
+        )
+        self._stop_event = threading.Event()
+        self._drained = threading.Event()
+        self.conn: Optional[psycopg.Connection] = None
+
+        self._worker_thread = threading.Thread(
+            target=self._worker, name="db-writer", daemon=True
+        )
+        self._worker_thread.start()
+
+    # -- Producer side (called from the MQTT callback thread) -----------
+
+    def write(self, topic: str, payload: Any) -> None:
+        """Enqueue a message for persistence; never blocks the caller."""
+        if self._stop_event.is_set():
+            logger.error("Writer stopped; dropping message on %s", topic)
+            return
+        try:
+            self._queue.put_nowait((datetime.now(timezone.utc), topic, payload))
+        except queue.Full:
+            # Make backpressure visible instead of stalling the MQTT loop.
+            logger.error(
+                "Write queue full (%d pending); dropping message on %s",
+                self._queue.maxsize,
+                topic,
+            )
+
+    # -- Consumer side (worker thread) ------------------------------------
+
+    def _worker(self) -> None:
+        while not self._stop_event.is_set() or not self._queue.empty():
+            batch = self._collect_batch()
+            if batch:
+                self._flush_with_retry(batch)
+        self._drained.set()
+        self._close_connection()
+
+    def _collect_batch(self) -> list[tuple[datetime, str, Any]]:
+        """Block for up to one flush interval, then drain up to batch size."""
+        batch: list[tuple[datetime, str, Any]] = []
+        try:
+            batch.append(self._queue.get(timeout=self._flush_interval))
+        except queue.Empty:
+            return batch
+        while len(batch) < self._batch_size:
+            try:
+                batch.append(self._queue.get_nowait())
+            except queue.Empty:
+                break
+        return batch
+
+    def _flush_with_retry(self, batch: list[tuple[datetime, str, Any]]) -> None:
+        for attempt in (1, 2):
+            try:
+                self._flush(batch)
+                return
+            except Exception:  # noqa: BLE001 - any DB error is retryable
+                logger.exception(
+                    "Failed to flush %d rows (attempt %d/2)", len(batch), attempt
+                )
+                self._close_connection()
+                if attempt == 1:
+                    time.sleep(2.0)
+
+        if self._stop_event.is_set():
+            logger.error("Shutting down; dropping %d unwritten rows", len(batch))
+            return
+        # Keep the data; it is retried on the next loop iteration.
+        logger.error("Re-enqueueing %d rows after failed flush", len(batch))
+        for item in batch:
+            try:
+                self._queue.put_nowait(item)
+            except queue.Full:
+                logger.error("Write queue full while re-enqueueing; dropping row")
+
+    def _flush(self, batch: list[tuple[datetime, str, Any]]) -> None:
+        self.ensure_connected()
+        assert self.conn is not None
+        rows = [
+            (ts, topic, self._device_id(payload), Jsonb(payload))
+            for ts, topic, payload in batch
+        ]
+        with self.conn.cursor() as cur, self.conn.transaction():
+            cur.executemany(_INSERT_SQL, rows)
+        logger.debug("Flushed %d rows to TimescaleDB", len(rows))
+
+    # -- Connection management ---------------------------------------------
+
+    def connect(self) -> None:
+        self.conn = psycopg.connect(
+            host=os.getenv("DB_HOST", "timescaledb"),
+            port=int(os.getenv("DB_PORT", "5432")),
+            dbname=os.getenv("DB_NAME", "telemetry"),
+            user=os.getenv("DB_USER", "ingester"),
+            password=os.getenv("DB_PASSWORD", "changeme"),
+            autocommit=True,
+        )
+        logger.info("Connected to TimescaleDB")
+
+    def ensure_connected(self) -> None:
+        if self.conn is None or self.conn.closed:
+            self.connect()
+            return
+        try:
+            # Cheap liveness check: a silently-dead socket (DB restart,
+            # network blip) would otherwise only surface on the next write.
+            self.conn.execute("SELECT 1").fetchone()
+        except psycopg.Error:
+            logger.warning("Database connection stale; reconnecting")
+            self._close_connection()
+            self.connect()
+
+    def _close_connection(self) -> None:
+        if self.conn is not None:
+            try:
+                self.conn.close()
+            except Exception:  # noqa: BLE001 - best-effort cleanup
+                pass
+            self.conn = None
+
+    # -- Lifecycle -----------------------------------------------------------
+
+    def close(self, timeout: float = 15.0) -> None:
+        """Stop the worker thread, flushing any queued rows first."""
+        if self._stop_event.is_set():
+            return
+        self._stop_event.set()
+        if not self._drained.wait(timeout):
+            logger.warning("Timed out flushing write queue; some rows may be lost")
+        self._close_connection()
+
+    @staticmethod
+    def _device_id(payload: Any) -> Optional[str]:
+        if isinstance(payload, dict):
+            return payload.get("device_id") or payload.get("deviceId")
+        return None

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

mqtt-client/main.py

@@ -0,0 +1,177 @@
+"""MQTT client that subscribes to an MQTTS broker and ingests telemetry.
+
+Supports three broker authentication modes (see configure_tls):
+  - TLS-PSK: symmetric pre-shared key, no certificates needed.
+  - Client certificate (mTLS): presented when the cert files exist.
+  - MQTT username/password: sent in the CONNECT packet; combinable with
+    either TLS mode.
+
+Persists received messages to TimescaleDB and logs locally and/or to Loki.
+"""
+
+import json
+import os
+import ssl
+import time
+
+import paho.mqtt.client as mqtt
+from paho.mqtt.enums import CallbackAPIVersion
+
+from db_writer import DbWriter
+from logging_setup import setup_logging
+
+logger = setup_logging()
+
+BROKER_HOST = os.getenv("MQTT_BROKER_HOST")
+BROKER_PORT = int(os.getenv("MQTT_BROKER_PORT", "8883"))
+TOPIC = os.getenv("MQTT_TOPIC", "sensors/#")
+CLIENT_ID = os.getenv("MQTT_CLIENT_ID", "mqtt-ingester")
+
+CA_CERT = os.getenv("MQTT_CA_CERT", "/certs/ca.crt")
+CLIENT_CERT = os.getenv("MQTT_CLIENT_CERT", "/certs/client.crt")
+CLIENT_KEY = os.getenv("MQTT_CLIENT_KEY", "/certs/client.key")
+# Empty string would be treated as a real (wrong) password by OpenSSL.
+CLIENT_KEY_PASSWORD = os.getenv("MQTT_CLIENT_KEY_PASSWORD") or None
+
+PSK_IDENTITY = os.getenv("MQTT_PSK_IDENTITY", "")
+PSK_HEX = os.getenv("MQTT_PSK_HEX", "")
+
+MQTT_USERNAME = os.getenv("MQTT_USERNAME", "")
+MQTT_PASSWORD = os.getenv("MQTT_PASSWORD", "")
+
+writer = DbWriter()
+
+
+def on_connect(client, userdata, flags, reason_code, properties):
+    if reason_code == 0:
+        logger.info("Connected to broker %s:%s", BROKER_HOST, BROKER_PORT)
+        client.subscribe(TOPIC, qos=1)
+        logger.info("Subscribed to topic %s", TOPIC)
+    else:
+        logger.error("Connection failed with reason code %s", reason_code)
+
+
+def on_message(client, userdata, msg):
+    topic = msg.topic
+    try:
+        payload = json.loads(msg.payload.decode("utf-8", errors="replace"))
+    except json.JSONDecodeError:
+        payload = {"raw": msg.payload.decode("utf-8", errors="replace")}
+
+    # Debug level: at INFO this logs (and possibly ships to Loki) once per
+    # message, which dominates ingestion cost at high message rates.
+    logger.debug("Received message on %s", topic, extra={"topic": topic})
+
+    # Enqueues only; all blocking DB work happens on the writer thread.
+    writer.write(topic, payload)
+
+
+def on_disconnect(client, userdata, reason_code, properties):
+    logger.warning("Disconnected from broker (reason %s)", reason_code)
+
+
+def configure_tls() -> ssl.SSLContext:
+    """Build the SSL context for the broker connection.
+
+    TLS-PSK takes precedence when MQTT_PSK_IDENTITY/MQTT_PSK_HEX are set;
+    otherwise a certificate-based context is used. The client certificate
+    is optional in the latter case, so brokers that authenticate via MQTT
+    username/password (or anonymous) also work.
+    """
+    if PSK_IDENTITY or PSK_HEX:
+        if not (PSK_IDENTITY and PSK_HEX):
+            raise RuntimeError(
+                "TLS-PSK requires both MQTT_PSK_IDENTITY and MQTT_PSK_HEX to be set"
+            )
+        if not getattr(ssl, "HAS_PSK", False):
+            raise RuntimeError(
+                "TLS-PSK requested but this Python/OpenSSL build has no PSK "
+                "support (ssl.HAS_PSK is False; requires Python 3.13+)"
+            )
+        try:
+            psk = bytes.fromhex(PSK_HEX)
+        except ValueError:
+            raise RuntimeError("MQTT_PSK_HEX is not a valid hex string") from None
+        if not psk:
+            raise RuntimeError("MQTT_PSK_HEX must not be empty")
+
+        context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+        # With PSK there is no certificate to verify; the shared key
+        # authenticates both sides. PSK cipher suites are TLS 1.2 only.
+        context.check_hostname = False
+        context.verify_mode = ssl.CERT_NONE
+        context.maximum_version = ssl.TLSVersion.TLSv1_2
+        context.set_ciphers("PSK")
+        context.set_psk_client_callback(lambda hint: (PSK_IDENTITY, psk))
+        logger.info("Using TLS-PSK authentication (identity %s)", PSK_IDENTITY)
+        return context
+
+    context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+    context.load_verify_locations(cafile=CA_CERT)
+    if os.path.exists(CLIENT_CERT) and os.path.exists(CLIENT_KEY):
+        context.load_cert_chain(CLIENT_CERT, CLIENT_KEY, CLIENT_KEY_PASSWORD)
+        logger.info("Using client certificate %s for broker auth", CLIENT_CERT)
+    else:
+        logger.info(
+            "Client certificate not found at %s; connecting with CA verification"
+            " only (broker must accept username/password or anonymous auth)",
+            CLIENT_CERT,
+        )
+    return context
+
+
+def main() -> None:
+    if not BROKER_HOST:
+        logger.error("MQTT_BROKER_HOST is not set")
+        return
+
+    client = mqtt.Client(
+        callback_api_version=CallbackAPIVersion.VERSION2,
+        client_id=CLIENT_ID,
+        protocol=mqtt.MQTTv5,
+    )
+
+    # TLS is always enabled: the stack is MQTTS-only. The mode (PSK vs
+    # certificates) and the presence of a client certificate are decided
+    # by configure_tls().
+    client.tls_set_context(configure_tls())
+
+    if MQTT_USERNAME:
+        # Credentials travel inside the CONNECT packet, so they work with
+        # any TLS mode (PSK, cert, or CA-only) and any broker config such
+        # as Mosquitto's password_file.
+        client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD)
+        logger.info("Using MQTT username/password authentication")
+    elif not PSK_IDENTITY and not (
+        os.path.exists(CLIENT_CERT) and os.path.exists(CLIENT_KEY)
+    ):
+        logger.warning(
+            "No broker authentication configured: no username/password, no PSK, "
+            "and no client certificate found at %s. The broker must allow "
+            "anonymous connections.",
+            CLIENT_CERT,
+        )
+
+    client.on_connect = on_connect
+    client.on_message = on_message
+    client.on_disconnect = on_disconnect
+
+    # loop_forever() reconnects automatically; this bounds the backoff so a
+    # long outage does not stretch into multi-minute gaps between attempts.
+    client.reconnect_delay_set(min_delay=1, max_delay=60)
+
+    try:
+        while True:
+            try:
+                client.connect(BROKER_HOST, BROKER_PORT, keepalive=60)
+                client.loop_forever()
+            except Exception:  # noqa: BLE001
+                logger.exception("Connection error, retrying in 5s")
+                time.sleep(5)
+    finally:
+        # Blocks briefly to flush any queued rows before the process exits.
+        writer.close()
+
+
+if __name__ == "__main__":
+    main()

mqtt-client/requirements.txt

@@ -0,0 +1,4 @@
+paho-mqtt==2.1.0
+psycopg[binary]==3.2.3
+python-json-logger==2.0.7
+requests==2.32.3