118 lines
4.5 KiB
Python
118 lines
4.5 KiB
Python
|
"""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
|