commit c2ae24818e7eeafec0be6d92e57e1e802b1bb95b Author: Jonas Bengtsson Date: Sun Aug 30 09:56:33 2026 +0200 🎉 Initial commit diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..2a43c40 --- /dev/null +++ b/.env.example @@ -0,0 +1,40 @@ +# --- Database --- +POSTGRES_USER=ingester +POSTGRES_PASSWORD=changeme +POSTGRES_DB=telemetry +POSTGRES_PORT=5432 + +# --- MQTT Broker (MQTTS) --- +MQTT_BROKER_HOST=broker.example.com +MQTT_BROKER_PORT=8883 +MQTT_TOPIC=sensors/# +MQTT_CLIENT_ID=mqtt-ingester + +# --- Broker authentication (all optional; combinable) --- +# MQTT username/password (CONNECT-packet auth, Mosquitto password_file etc.). +MQTT_USERNAME= +MQTT_PASSWORD= +# TLS-PSK: takes precedence over certificates when both are set. +# MQTT_PSK_HEX is a hex string, e.g. "0123456789abcdef" (no 0x prefix). +MQTT_PSK_IDENTITY= +MQTT_PSK_HEX= + +# --- Batched DB writes (mqtt-client) --- +DB_BATCH_SIZE=500 +DB_FLUSH_SECONDS=1.0 + +# --- Logging --- +LOG_LEVEL=INFO +LOG_TO_FILE=true + +# --- Loki (optional, leave LOKI_URL empty to disable) --- +LOKI_URL= +LOKI_USER= +LOKI_PASSWORD= +LOKI_FLUSH_SECONDS=5 +LOKI_TIMEOUT_SECONDS=5 + +# --- API --- +API_PORT=8443 +DB_POOL_MIN_SIZE=2 +DB_POOL_MAX_SIZE=10 diff --git a/README.md b/README.md new file mode 100644 index 0000000..8943686 --- /dev/null +++ b/README.md @@ -0,0 +1,205 @@ +# MQTT Ingester Stack + +A Docker-based stack that ingests time-series data from an MQTTS broker into +TimescaleDB, exposes it through an authenticated HTTPS API, and optionally +forwards logs to Loki. + +## Architecture + +```mermaid +flowchart LR + Broker[MQTTS Broker] -->|TLS + client cert| MQTT[MQTT Client] + MQTT -->|SQL| TSDB[(TimescaleDB)] + MQTT -->|logs| Loki[Loki] + TSDB --> API[HTTPS API] + Client[API Client] -->|mTLS| API +``` + +| Service | Purpose | +| ------------ | -------------------------------------------------------------- | +| `timescaledb`| Time-series database storing telemetry as hypertables. | +| `mqtt-client`| Subscribes to MQTTS topics, writes payloads to the database. | +| `api` | FastAPI service serving telemetry over HTTPS with mTLS. | +| `loki` | Optional log aggregation (enabled via the `telemetry` profile).| + +## Prerequisites + +- Docker and Docker Compose v2. +- A CA certificate and signed certificates for the MQTT client, the API + server, and any API clients. + +## Configuration + +Copy the example environment file and edit it: + +```bash +cp .env.example .env +``` + +Key variables: + +| Variable | Description | +| ------------------- | --------------------------------------------- | +| `POSTGRES_USER` | TimescaleDB user. | +| `POSTGRES_PASSWORD` | TimescaleDB password. | +| `MQTT_BROKER_HOST` | MQTTS broker hostname. | +| `MQTT_BROKER_PORT` | MQTTS broker port (default `8883`). | +| `MQTT_TOPIC` | Topic(s) to subscribe to (default `sensors/#`).| +| `MQTT_USERNAME` | MQTT username for broker auth (optional). | +| `MQTT_PASSWORD` | MQTT password for broker auth (optional). | +| `MQTT_PSK_IDENTITY` | TLS-PSK identity (optional; enables PSK mode).| +| `MQTT_PSK_HEX` | TLS-PSK key as hex (optional). | +| `DB_BATCH_SIZE` | Rows per batched DB insert (default `500`). | +| `DB_FLUSH_SECONDS` | Max seconds a batch waits before flushing (default `1.0`). | +| `DB_POOL_MIN_SIZE` | Min API DB pool connections (default `2`). | +| `DB_POOL_MAX_SIZE` | Max API DB pool connections (default `10`). | +| `LOKI_URL` | Loki push URL; leave empty to disable. | +| `LOKI_FLUSH_SECONDS`| Seconds between batched Loki pushes (default `5`). | +| `API_PORT` | Host port for the HTTPS API (default `8443`). | + +## Certificates + +Place the following files in the `certs/` directory: + +| File | Used by | Purpose | +| ------------ | ------------ | -------------------------------- | +| `ca.crt` | all | CA certificate. | +| `client.crt` | mqtt-client | Client certificate for the broker.| +| `client.key` | mqtt-client | Client private key. | +| `server.crt` | api | API server certificate. | +| `server.key` | api | API server private key. | + +Generate a self-signed CA and certificates for testing: + +```bash +mkdir -p certs && cd certs + +# CA +openssl req -x509 -newkey rsa:2048 -days 365 -nodes \ + -keyout ca.key -out ca.crt -subj "/CN=Test CA" + +# MQTT client certificate +openssl req -newkey rsa:2048 -nodes -keyout client.key -out client.csr \ + -subj "/CN=mqtt-client" +openssl x509 -req -in client.csr -CA ca.crt -CAkey ca.key \ + -CAcreateserial -out client.crt -days 365 + +# API server certificate +openssl req -newkey rsa:2048 -nodes -keyout server.key -out server.csr \ + -subj "/CN=localhost" +openssl x509 -req -in server.csr -CA ca.crt -CAkey ca.key \ + -CAcreateserial -out server.crt -days 365 +``` + +> The broker must trust `ca.crt` and accept `client.crt` for authentication. + +## MQTT Broker Authentication + +The client always connects over TLS and picks its authentication mode from +the environment: + +| Mode | Variables | Notes | +| ---- | --------- | ----- | +| **TLS-PSK** | `MQTT_PSK_IDENTITY` + `MQTT_PSK_HEX` | Takes precedence when set. No certificates needed; the broker must have a PSK listener (Mosquitto: `psk_hint` + `psk_file`). Requires Python 3.13+ (the mqtt-client image provides this). | +| **Client certificate (mTLS)** | presence of `certs/client.crt` + `certs/client.key` | Default when the files exist and PSK is not configured. | +| **Username/password** | `MQTT_USERNAME` (+ `MQTT_PASSWORD`) | Sent in the MQTT CONNECT packet; combinable with either TLS mode. Matches Mosquitto's `password_file`. | +| **Anonymous** | none of the above | Only works if the broker allows it; a warning is logged. | + +Examples: + +```bash +# TLS-PSK +MQTT_PSK_IDENTITY=ingester MQTT_PSK_HEX=0123456789abcdef + +# Username/password over CA-verified TLS (no client cert required) +MQTT_USERNAME=ingester MQTT_PASSWORD=secret +``` + +TLS-PSK in Mosquitto requires a dedicated listener, e.g.: + +``` +listener 8884 +psk_hint mqtt-ingester +psk_file /mosquitto/psk_file # line: ingester:0123456789abcdef +``` + +## Running + +Start the core stack: + +```bash +docker compose up -d +``` + +Start with Loki enabled: + +```bash +docker compose --profile telemetry up -d +``` + +## API + +The API is served over HTTPS with mutual TLS. Clients must present a +certificate signed by the configured CA. + +```bash +curl --cacert certs/ca.crt \ + --cert certs/client.crt \ + --key certs/client.key \ + https://localhost:8443/health +``` + +Endpoints: + +| Method | Path | Description | +| ------ | -------------------- | ---------------------------------------- | +| GET | `/health` | Health check. | +| GET | `/telemetry` | List telemetry (filters: `topic`, `device_id`, `limit`). | +| GET | `/telemetry/latest` | Latest record (filters: `topic`, `device_id`). | + +Example: + +```bash +curl --cacert certs/ca.crt --cert certs/client.crt --key certs/client.key \ + "https://localhost:8443/telemetry?topic=sensors/temp&limit=10" +``` + +## Logging + +The MQTT client logs in JSON format to: + +- **Console** (always). +- **File** (`/var/log/mqtt-client/mqtt-client.log`) when `LOG_TO_FILE=true`. +- **Loki** when `LOKI_URL` is set. Records are buffered and pushed in + batches from a background thread, so a slow or unavailable Loki never + blocks message ingestion. + +## Performance + +- The MQTT client buffers incoming messages in an in-memory queue and + inserts them into TimescaleDB in batches (one transaction per batch), + keeping database I/O off the MQTT network loop. Tune with + `DB_BATCH_SIZE` and `DB_FLUSH_SECONDS`. If the queue fills (broker + outpacing the database), the client logs the condition and drops new + messages rather than stalling the broker connection. +- The API serves queries from a pool of reused database connections + (`DB_POOL_MIN_SIZE`/`DB_POOL_MAX_SIZE`). +- Chunks older than 7 days are compressed, and data older than 365 days + is dropped automatically (see `db/init/001_schema.sql`). + +## Database Schema + +Telemetry is stored in the `telemetry` hypertable: + +| Column | Type | Description | +| ---------- | ----------- | ------------------------------ | +| `time` | `TIMESTAMPTZ`| Ingestion timestamp. | +| `topic` | `TEXT` | MQTT topic. | +| `device_id`| `TEXT` | Extracted from payload if present.| +| `payload` | `JSONB` | Full message payload. | + +## Security Notes + +- Change `POSTGRES_PASSWORD` before deploying. +- Use a real CA and rotate certificates regularly. +- Restrict the API port with a firewall or reverse proxy. diff --git a/api/.dockerignore b/api/.dockerignore new file mode 100644 index 0000000..9068aa1 --- /dev/null +++ b/api/.dockerignore @@ -0,0 +1,5 @@ +__pycache__/ +*.pyc +*.log +.env +.venv/ \ No newline at end of file diff --git a/api/Dockerfile b/api/Dockerfile new file mode 100644 index 0000000..812dc98 --- /dev/null +++ b/api/Dockerfile @@ -0,0 +1,13 @@ +FROM python:3.12-slim + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY . . + +RUN useradd --system --no-create-home api +USER api + +CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8443", "--ssl-keyfile", "/certs/server.key", "--ssl-certfile", "/certs/server.crt", "--ssl-ca-certs", "/certs/ca.crt"] diff --git a/api/main.py b/api/main.py new file mode 100644 index 0000000..899d629 --- /dev/null +++ b/api/main.py @@ -0,0 +1,132 @@ +"""FastAPI service exposing telemetry data from TimescaleDB over HTTPS. + +Requires mutual TLS: clients must present a certificate signed by the CA. +Requests are served from a pool of database connections, so each request +reuses an established session instead of paying connect/auth costs. +""" + +import os +from contextlib import asynccontextmanager +from typing import Any, Final, Optional + +import psycopg +from fastapi import FastAPI, HTTPException, Query, Request +from fastapi.responses import JSONResponse +from psycopg import sql +from psycopg.rows import dict_row +from psycopg_pool import ConnectionPool + +_SELECT_SQL: Final = "SELECT time, topic, device_id, payload FROM telemetry" + +# Connections use row_factory=dict_row, so each row is a dict. +_DictConnection = psycopg.Connection[dict[str, Any]] + +pool: Optional[ConnectionPool[_DictConnection]] = None + + +def _connect_kwargs() -> dict[str, Any]: + return { + "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"), + "row_factory": dict_row, + } + + +@asynccontextmanager +async def lifespan(app: FastAPI): + global pool + # Subscripting solves the pool's invariant connection type; row_factory + # in kwargs makes every pooled connection return dict rows. + server_pool = ConnectionPool[_DictConnection]( + kwargs=_connect_kwargs(), + min_size=int(os.getenv("DB_POOL_MIN_SIZE", "2")), + max_size=int(os.getenv("DB_POOL_MAX_SIZE", "10")), + ) + server_pool.open(wait=False) + pool = server_pool + try: + yield + finally: + pool = None + server_pool.close() + + +app = FastAPI(title="MQTT Telemetry API", lifespan=lifespan) + + +def get_pool() -> ConnectionPool[_DictConnection]: + if pool is None: + raise RuntimeError("Database pool not initialized") + return pool + + +@app.middleware("http") +async def require_client_cert(request: Request, call_next): + # uvicorn populates the client cert into the transport when mTLS is enabled. + cert = getattr(request.scope.get("transport"), "get_extra_info", lambda _: None)( + "peercert" + ) + if cert is None: + # Return a response directly: HTTPException raised from middleware is + # not handled by the (inner) exception handlers and would surface + # as a 500. + return JSONResponse(status_code=401, content={"detail": "Client certificate required"}) + return await call_next(request) + + +@app.get("/health") +def health() -> dict[str, str]: + return {"status": "ok"} + + +def _telemetry_query( + topic: Optional[str], device_id: Optional[str], limit: int +) -> tuple[sql.Composed, list[Any]]: + """Build the SELECT for the requested filters as composed SQL. + + The statement is composed with psycopg.sql, and all values are bound + as parameters, so nothing is interpolated into the query string. + """ + conditions: list[sql.Composable] = [] + params: list[Any] = [] + if topic: + conditions.append(sql.SQL("topic = %s")) + params.append(topic) + if device_id: + conditions.append(sql.SQL("device_id = %s")) + params.append(device_id) + query = sql.SQL(_SELECT_SQL) + if conditions: + query += sql.SQL(" WHERE ") + sql.SQL(" AND ").join(conditions) + query += sql.SQL(" ORDER BY time DESC LIMIT %s") + params.append(limit) + return query, params + + +@app.get("/telemetry") +def list_telemetry( + topic: Optional[str] = Query(default=None), + device_id: Optional[str] = Query(default=None), + limit: int = Query(default=100, le=1000), +) -> list[dict[str, Any]]: + query, params = _telemetry_query(topic, device_id, limit) + with get_pool().connection() as conn, conn.cursor() as cur: + cur.execute(query, params) + return cur.fetchall() + + +@app.get("/telemetry/latest") +def latest_telemetry( + topic: Optional[str] = Query(default=None), + device_id: Optional[str] = Query(default=None), +) -> Optional[dict[str, Any]]: + query, params = _telemetry_query(topic, device_id, 1) + with get_pool().connection() as conn, conn.cursor() as cur: + cur.execute(query, params) + row = cur.fetchone() + if row is None: + raise HTTPException(status_code=404, detail="No data found") + return row diff --git a/api/requirements.txt b/api/requirements.txt new file mode 100644 index 0000000..0408c39 --- /dev/null +++ b/api/requirements.txt @@ -0,0 +1,3 @@ +fastapi==0.115.6 +uvicorn[standard]==0.34.0 +psycopg[binary,pool]==3.2.3 diff --git a/db/init/001_schema.sql b/db/init/001_schema.sql new file mode 100644 index 0000000..d39ccb2 --- /dev/null +++ b/db/init/001_schema.sql @@ -0,0 +1,32 @@ +-- TimescaleDB schema for MQTT telemetry ingestion. +-- Executed automatically on first container start. + +CREATE EXTENSION IF NOT EXISTS timescaledb; + +CREATE TABLE IF NOT EXISTS telemetry ( + time TIMESTAMPTZ NOT NULL, + topic TEXT NOT NULL, + device_id TEXT, + payload JSONB NOT NULL +); + +-- Convert to hypertable for time-series storage. +SELECT create_hypertable('telemetry', 'time', if_not_exists => TRUE); + +-- Indexes for common query patterns. +-- (topic, time) and (device_id, time) serve the filtered API queries; +-- a bare (time) index serves the unfiltered /telemetry/latest lookup, +-- which would otherwise require a full backward scan of the hypertable. +CREATE INDEX IF NOT EXISTS idx_telemetry_topic_time ON telemetry (topic, time DESC); +CREATE INDEX IF NOT EXISTS idx_telemetry_device_time ON telemetry (device_id, time DESC); +CREATE INDEX IF NOT EXISTS idx_telemetry_time ON telemetry (time DESC); + +-- Compression: older chunks are squeezed on a background schedule. +ALTER TABLE telemetry SET ( + timescaledb.compress, + timescaledb.compress_segmentby = 'topic, device_id' +); + +-- Compress chunks older than 7 days, drop data older than 1 year. +SELECT add_compression_policy('telemetry', INTERVAL '7 days'); +SELECT add_retention_policy('telemetry', INTERVAL '365 days'); diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..f9ec6f1 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,98 @@ +services: + timescaledb: + image: timescale/timescaledb:2.17.2-pg16 + container_name: mqtt-timescaledb + restart: unless-stopped + environment: + POSTGRES_USER: ${POSTGRES_USER:-ingester} + POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-changeme} + POSTGRES_DB: ${POSTGRES_DB:-telemetry} + ports: + - "127.0.0.1:${POSTGRES_PORT:-5432}:5432" + volumes: + - timescale-data:/var/lib/postgresql/data + - ./db/init:/docker-entrypoint-initdb.d:ro + healthcheck: + test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER:-ingester} -d ${POSTGRES_DB:-telemetry}"] + interval: 10s + timeout: 5s + retries: 5 + + mqtt-client: + build: ./mqtt-client + container_name: mqtt-client + restart: unless-stopped + depends_on: + timescaledb: + condition: service_healthy + environment: + MQTT_BROKER_HOST: ${MQTT_BROKER_HOST} + MQTT_BROKER_PORT: ${MQTT_BROKER_PORT:-8883} + MQTT_TOPIC: ${MQTT_TOPIC:-sensors/#} + MQTT_CLIENT_ID: ${MQTT_CLIENT_ID:-mqtt-ingester} + MQTT_CA_CERT: /certs/ca.crt + MQTT_CLIENT_CERT: /certs/client.crt + MQTT_CLIENT_KEY: /certs/client.key + MQTT_USERNAME: ${MQTT_USERNAME:-} + MQTT_PASSWORD: ${MQTT_PASSWORD:-} + MQTT_PSK_IDENTITY: ${MQTT_PSK_IDENTITY:-} + MQTT_PSK_HEX: ${MQTT_PSK_HEX:-} + DB_HOST: timescaledb + DB_PORT: 5432 + DB_NAME: ${POSTGRES_DB:-telemetry} + DB_USER: ${POSTGRES_USER:-ingester} + DB_PASSWORD: ${POSTGRES_PASSWORD:-changeme} + DB_BATCH_SIZE: ${DB_BATCH_SIZE:-500} + DB_FLUSH_SECONDS: ${DB_FLUSH_SECONDS:-1.0} + LOG_LEVEL: ${LOG_LEVEL:-INFO} + LOG_TO_FILE: ${LOG_TO_FILE:-true} + LOG_FILE_PATH: /var/log/mqtt-client/mqtt-client.log + LOKI_URL: ${LOKI_URL:-} + LOKI_USER: ${LOKI_USER:-} + LOKI_PASSWORD: ${LOKI_PASSWORD:-} + LOKI_FLUSH_SECONDS: ${LOKI_FLUSH_SECONDS:-5} + LOKI_TIMEOUT_SECONDS: ${LOKI_TIMEOUT_SECONDS:-5} + volumes: + - ./certs:/certs:ro + - mqtt-logs:/var/log/mqtt-client + + api: + build: ./api + container_name: mqtt-api + restart: unless-stopped + depends_on: + timescaledb: + condition: service_healthy + environment: + DB_HOST: timescaledb + DB_PORT: 5432 + DB_NAME: ${POSTGRES_DB:-telemetry} + DB_USER: ${POSTGRES_USER:-ingester} + DB_PASSWORD: ${POSTGRES_PASSWORD:-changeme} + DB_POOL_MIN_SIZE: ${DB_POOL_MIN_SIZE:-2} + DB_POOL_MAX_SIZE: ${DB_POOL_MAX_SIZE:-10} + API_HOST: 0.0.0.0 + API_PORT: 8443 + API_CA_CERT: /certs/ca.crt + API_SERVER_CERT: /certs/server.crt + API_SERVER_KEY: /certs/server.key + ports: + - "${API_PORT:-8443}:8443" + volumes: + - ./certs:/certs:ro + + loki: + image: grafana/loki:3.4.2 + container_name: mqtt-loki + restart: unless-stopped + profiles: ["telemetry"] + command: -config.file=/etc/loki/local-config.yaml + ports: + - "127.0.0.1:3100:3100" + volumes: + - loki-data:/loki + +volumes: + timescale-data: + mqtt-logs: + loki-data: diff --git a/mqtt-client/.dockerignore b/mqtt-client/.dockerignore new file mode 100644 index 0000000..9068aa1 --- /dev/null +++ b/mqtt-client/.dockerignore @@ -0,0 +1,5 @@ +__pycache__/ +*.pyc +*.log +.env +.venv/ \ No newline at end of file diff --git a/mqtt-client/Dockerfile b/mqtt-client/Dockerfile new file mode 100644 index 0000000..6780ce5 --- /dev/null +++ b/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"] diff --git a/mqtt-client/db_writer.py b/mqtt-client/db_writer.py new file mode 100644 index 0000000..1a6d8f4 --- /dev/null +++ b/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 diff --git a/mqtt-client/logging_setup.py b/mqtt-client/logging_setup.py new file mode 100644 index 0000000..a6d0f75 --- /dev/null +++ b/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 diff --git a/mqtt-client/main.py b/mqtt-client/main.py new file mode 100644 index 0000000..edc44bc --- /dev/null +++ b/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() diff --git a/mqtt-client/requirements.txt b/mqtt-client/requirements.txt new file mode 100644 index 0000000..2385f9b --- /dev/null +++ b/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