155 lines
5.4 KiB
Python
155 lines
5.4 KiB
Python
"""Collector Manager - orchestrates all device collectors."""
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import select
|
|
|
|
from app.core.config import get_settings
|
|
from app.core.database import async_session
|
|
from app.models.device import Device
|
|
from app.collectors.base import BaseCollector
|
|
from app.collectors.modbus_tcp import ModbusTcpCollector
|
|
from app.collectors.mqtt_collector import MqttCollector
|
|
from app.collectors.http_collector import HttpCollector
|
|
|
|
logger = logging.getLogger("collector.manager")
|
|
|
|
# Full registry mapping protocol names to collector classes
|
|
COLLECTOR_REGISTRY: dict[str, type[BaseCollector]] = {
|
|
"modbus_tcp": ModbusTcpCollector,
|
|
"mqtt": MqttCollector,
|
|
"http_api": HttpCollector,
|
|
}
|
|
|
|
|
|
def get_enabled_collectors() -> dict[str, type[BaseCollector]]:
|
|
"""Return collector registry filtered by customer config.
|
|
|
|
If the customer config specifies a 'collectors' list, only those
|
|
protocols are enabled. Otherwise fall back to the full registry.
|
|
"""
|
|
settings = get_settings()
|
|
customer_config = settings.load_customer_config()
|
|
enabled_list = customer_config.get("collectors")
|
|
if enabled_list is None:
|
|
return COLLECTOR_REGISTRY
|
|
enabled = {}
|
|
for name in enabled_list:
|
|
if name in COLLECTOR_REGISTRY:
|
|
enabled[name] = COLLECTOR_REGISTRY[name]
|
|
else:
|
|
logger.warning("Customer config references unknown collector '%s', skipping", name)
|
|
return enabled
|
|
|
|
|
|
class CollectorManager:
|
|
"""Manages lifecycle of all device collectors."""
|
|
|
|
def __init__(self):
|
|
self._collectors: dict[int, BaseCollector] = {} # device_id -> collector
|
|
self._running = False
|
|
|
|
async def start(self):
|
|
"""Load active devices from DB and start their collectors."""
|
|
self._running = True
|
|
await self._load_and_start_collectors()
|
|
logger.info("CollectorManager started with %d collectors", len(self._collectors))
|
|
|
|
async def stop(self):
|
|
"""Stop all collectors."""
|
|
self._running = False
|
|
for device_id, collector in self._collectors.items():
|
|
try:
|
|
await collector.stop()
|
|
except Exception as e:
|
|
logger.error("Error stopping collector for device %d: %s", device_id, e)
|
|
self._collectors.clear()
|
|
logger.info("CollectorManager stopped")
|
|
|
|
async def _load_and_start_collectors(self):
|
|
"""Load active devices with supported protocols and start collectors."""
|
|
enabled = get_enabled_collectors()
|
|
logger.info("Enabled collectors: %s", list(enabled.keys()))
|
|
async with async_session() as session:
|
|
result = await session.execute(
|
|
select(Device).where(
|
|
Device.is_active == True,
|
|
Device.protocol.in_(list(enabled.keys())),
|
|
)
|
|
)
|
|
devices = result.scalars().all()
|
|
|
|
for device in devices:
|
|
await self.start_collector(
|
|
device.id,
|
|
device.code,
|
|
device.protocol,
|
|
device.connection_params or {},
|
|
device.collect_interval or 15,
|
|
)
|
|
|
|
async def start_collector(
|
|
self,
|
|
device_id: int,
|
|
device_code: str,
|
|
protocol: str,
|
|
connection_params: dict,
|
|
collect_interval: int,
|
|
) -> bool:
|
|
"""Start a single collector for a device."""
|
|
if device_id in self._collectors:
|
|
logger.warning("Collector already running for device %d", device_id)
|
|
return False
|
|
|
|
collector_cls = COLLECTOR_REGISTRY.get(protocol)
|
|
if not collector_cls:
|
|
logger.warning("No collector for protocol '%s' (device %s)", protocol, device_code)
|
|
return False
|
|
|
|
collector = collector_cls(device_id, device_code, connection_params, collect_interval)
|
|
self._collectors[device_id] = collector
|
|
await collector.start()
|
|
logger.info("Started %s collector for %s", protocol, device_code)
|
|
return True
|
|
|
|
async def stop_collector(self, device_id: int) -> bool:
|
|
"""Stop collector for a specific device."""
|
|
collector = self._collectors.pop(device_id, None)
|
|
if not collector:
|
|
return False
|
|
await collector.stop()
|
|
return True
|
|
|
|
async def restart_collector(self, device_id: int) -> bool:
|
|
"""Restart collector for a device by reloading its config from DB."""
|
|
await self.stop_collector(device_id)
|
|
async with async_session() as session:
|
|
result = await session.execute(
|
|
select(Device).where(Device.id == device_id)
|
|
)
|
|
device = result.scalar_one_or_none()
|
|
if not device or not device.is_active:
|
|
return False
|
|
return await self.start_collector(
|
|
device.id,
|
|
device.code,
|
|
device.protocol,
|
|
device.connection_params or {},
|
|
device.collect_interval or 15,
|
|
)
|
|
|
|
def get_collector(self, device_id: int) -> Optional[BaseCollector]:
|
|
return self._collectors.get(device_id)
|
|
|
|
def get_all_status(self) -> list[dict]:
|
|
"""Return status of all collectors."""
|
|
return [c.get_status() for c in self._collectors.values()]
|
|
|
|
@property
|
|
def collector_count(self) -> int:
|
|
return len(self._collectors)
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
return self._running
|