🎉 Initial commit

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