ohmu/dev/varvusensor/collector/collector.py

"""AS7341 värvianduri koguja: MQTT -> ClickHouse (uus67.varvus).

Tellib MQTT-teemat (vaikimisi varvus/spectral), kus ESPHome saadab
JSON-i kanalitega f1..f8, clear, nir (+ gain). Iga sonum = uks rida.
Lisaks arvutatakse ligikaudne värvus (CIE XYZ -> sRGB, vt color.py).

Tellimus tehakse on_connect-is, et see taastuks ka parast brokeri
taaskäivitust / ühenduse katkemist (clean session).
"""
from __future__ import annotations

import json
import os
from datetime import datetime
from zoneinfo import ZoneInfo

import paho.mqtt.client as mqtt
import requests

import color

MQTT_HOST = os.environ.get("MQTT_HOST", "mosquitto")
MQTT_PORT = int(os.environ.get("MQTT_PORT", "1883"))
MQTT_USER = os.environ.get("MQTT_USER", "")
MQTT_PASS = os.environ.get("MQTT_PASS", "")
MQTT_TOPIC = os.environ.get("MQTT_TOPIC", "varvus/spectral")
MQTT_TLS = os.environ.get("MQTT_TLS", "false").strip().lower() in ("1", "true", "yes", "on")

CH_URL = os.environ.get("CH_URL", "http://clickhouse:8123").rstrip("/")
CH_DB = os.environ.get("CH_DB", "uus67")
CH_TABLE = os.environ.get("CH_TABLE", "varvus")
CH_USER = os.environ.get("CH_USER", "zigbee")
CH_PASS = os.environ.get("CH_PASS", "")

TZ = ZoneInfo("Europe/Tallinn")
FIELDS = ("f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8", "clear", "nir")


def insert(row: dict) -> None:
    resp = requests.post(
        f"{CH_URL}/",
        params={
            "database": CH_DB,
            "query": f"INSERT INTO {CH_DB}.{CH_TABLE} FORMAT JSONEachRow",
        },
        data=json.dumps(row, ensure_ascii=False) + "\n",
        auth=(CH_USER, CH_PASS),
        timeout=10,
    )
    resp.raise_for_status()


def on_connect(client, userdata, flags, reason_code, properties=None) -> None:
    if reason_code == 0:
        client.subscribe(MQTT_TOPIC, qos=0)
        print(f"MQTT ühendatud, tellitud {MQTT_TOPIC}", flush=True)
    else:
        print(f"MQTT ühenduse viga: {reason_code}", flush=True)


def on_message(client, userdata, msg) -> None:
    try:
        payload = json.loads(msg.payload.decode())
    except Exception as exc:
        print(f"JSON viga: {exc}", flush=True)
        return
    row: dict = {
        "ts": datetime.now(TZ).strftime("%Y-%m-%d %H:%M:%S"),
        "device": "varvusensor",
    }
    channels: dict[str, float] = {}
    for field in FIELDS:
        value = payload.get(field)
        if value is not None:
            channels[field] = float(value)
            row[field] = float(value)
    gain = payload.get("gain")
    if gain is not None:
        row["gain"] = int(gain)
    row.update(color.analyze(channels))
    try:
        insert(row)
        print(
            f"salvestatud: {row['rgb_hex']} x={row['x_chrom']} y={row['y_chrom']} "
            f"f5={row.get('f5')} f6={row.get('f6')}",
            flush=True,
        )
    except Exception as exc:
        print(f"ClickHouse viga: {exc}", flush=True)


def main() -> None:
    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
    if MQTT_TLS:
        client.tls_set()
    if MQTT_USER:
        client.username_pw_set(MQTT_USER, MQTT_PASS)
    client.on_connect = on_connect
    client.on_message = on_message
    client.connect(MQTT_HOST, MQTT_PORT, 60)
    print(f"Koguja käivitub, broker {MQTT_HOST}:{MQTT_PORT} (TLS={MQTT_TLS})", flush=True)
    client.loop_forever()


if __name__ == "__main__":
    main()