Squashed 'core/' content from commit 92ec910

git-subtree-dir: core
git-subtree-split: 92ec910a132e379a3a6e442a75bcb07cac0f0010

backend/app/collectors/mqtt_collector.py

@@ -0,0 +1,117 @@
+"""MQTT protocol collector."""
+import json
+from typing import Optional
+
+import aiomqtt
+
+from app.collectors.base import BaseCollector
+
+
+class MqttCollector(BaseCollector):
+    """Collect data from devices via MQTT subscription.
+
+    connection_params example:
+    {
+        "broker": "localhost",
+        "port": 1883,
+        "topic": "device/INV-001/data",
+        "username": "",
+        "password": "",
+        "data_mapping": {
+            "active_power": {"key": "power", "unit": "kW"},
+            "voltage": {"key": "voltage", "unit": "V"}
+        }
+    }
+    """
+
+    def __init__(self, device_id, device_code, connection_params, collect_interval=15):
+        super().__init__(device_id, device_code, connection_params, collect_interval)
+        self._broker = connection_params.get("broker", "localhost")
+        self._port = connection_params.get("port", 1883)
+        self._topic = connection_params.get("topic", f"device/{device_code}/data")
+        self._username = connection_params.get("username", "") or None
+        self._password = connection_params.get("password", "") or None
+        self._data_mapping = connection_params.get("data_mapping", {})
+        self._client: Optional[aiomqtt.Client] = None
+        self._latest_data: dict = {}
+
+    async def connect(self):
+        # Connection is established in the run loop via context manager
+        pass
+
+    async def disconnect(self):
+        self._client = None
+
+    async def collect(self) -> dict:
+        # Return latest received data; cleared after read
+        data = self._latest_data.copy()
+        self._latest_data.clear()
+        return data
+
+    async def _run(self):
+        """Override run loop to use MQTT's push-based model."""
+        while self._running:
+            try:
+                async with aiomqtt.Client(
+                    self._broker,
+                    port=self._port,
+                    username=self._username,
+                    password=self._password,
+                ) as client:
+                    self._client = client
+                    self.status = "connected"
+                    self.last_error = None
+                    self._backoff = 1
+                    self.logger.info("MQTT connected to %s:%d", self._broker, self._port)
+
+                    await client.subscribe(self._topic)
+                    self.logger.info("Subscribed to %s", self._topic)
+
+                    async for message in client.messages:
+                        if not self._running:
+                            break
+                        try:
+                            payload = json.loads(message.payload.decode())
+                            data = self._parse_payload(payload)
+                            if data:
+                                self._latest_data.update(data)
+                                await self._save_data(data)
+                                from datetime import datetime, timezone
+                                self.last_collect_time = datetime.now(timezone.utc)
+                        except (json.JSONDecodeError, ValueError) as e:
+                            self.logger.warning("Bad MQTT payload on %s: %s", message.topic, e)
+
+            except aiomqtt.MqttError as e:
+                self.status = "error"
+                self.last_error = str(e)
+                self.logger.error("MQTT error for %s: %s", self.device_code, e)
+                await self._wait_backoff()
+            except Exception as e:
+                self.status = "error"
+                self.last_error = str(e)
+                self.logger.error("Unexpected MQTT error for %s: %s", self.device_code, e)
+                await self._wait_backoff()
+
+        self.status = "disconnected"
+
+    def _parse_payload(self, payload: dict) -> dict:
+        """Parse MQTT JSON payload into data points.
+
+        If data_mapping is configured, use it. Otherwise, treat all
+        numeric top-level keys as data points with empty units.
+        """
+        data = {}
+        if self._data_mapping:
+            for data_type, mapping in self._data_mapping.items():
+                key = mapping.get("key", data_type)
+                unit = mapping.get("unit", "")
+                if key in payload:
+                    try:
+                        data[data_type] = (float(payload[key]), unit)
+                    except (TypeError, ValueError):
+                        pass
+        else:
+            for key, value in payload.items():
+                if isinstance(value, (int, float)):
+                    data[key] = (float(value), "")
+        return data