186 lines
6.6 KiB
Python
186 lines
6.6 KiB
Python
"""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
|