PCM_Viewer/influxdb_wrapper.py

"""
InfluxDB 客户端模块
"""
from PyQt6.QtCore import QObject, pyqtSignal as Signal, QTimer, pyqtSlot as Slot
from influxdb_client import InfluxDBClient as InfluxClientLib
from typing import Dict, List, Optional
import logging

logger = logging.getLogger(__name__)


class InfluxDBClient(QObject):
    """
    InfluxDB 客户端，用于获取实时数据
    """
    dataReceived = Signal(dict)  # 数据接收信号，参数为 {field_name: value}
    errorOccurred = Signal(str)  # 错误信号
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self._client: Optional[InfluxClientLib] = None
        self._url: str = ""
        self._token: str = ""
        self._org: str = ""
        self._bucket: str = ""
        self._query: str = ""
        self._timer: Optional[QTimer] = None
        self._is_connected: bool = False
        
    @Slot(str, str, str, str)
    def connect(self, url: str, token: str, org: str, bucket: str):
        """连接到 InfluxDB"""
        try:
            self._url = url
            self._token = token
            self._org = org
            self._bucket = bucket
            
            if self._client:
                self._client.close()
            
            self._client = InfluxClientLib(
                url=url,
                token=token,
                org=org
            )
            
            # 测试连接
            query_api = self._client.query_api()
            test_query = f'from(bucket: "{bucket}") |> range(start: -1m) |> limit(n: 1)'
            query_api.query(test_query)
            
            self._is_connected = True
            logger.info(f"Connected to InfluxDB: {url}")
            
        except Exception as e:
            self._is_connected = False
            error_msg = f"连接失败: {str(e)}"
            logger.error(error_msg)
            self.errorOccurred.emit(error_msg)
    
    @Slot(str)
    def setQuery(self, query: str):
        """设置查询语句"""
        self._query = query
    
    @Slot(int)
    def startQuery(self, interval_ms: int = 1000):
        """开始定时查询，interval_ms 为毫秒"""
        if not self._is_connected:
            self.errorOccurred.emit("未连接到 InfluxDB")
            return
        
        if self._timer:
            self._timer.stop()
        
        self._timer = QTimer(self)
        self._timer.timeout.connect(self._executeQuery)
        self._timer.start(interval_ms)
        # 立即执行一次
        self._executeQuery()
    
    @Slot()
    def stopQuery(self):
        """停止查询"""
        if self._timer:
            self._timer.stop()
    
    def _executeQuery(self):
        """执行查询"""
        if not self._is_connected or not self._query:
            return
        
        try:
            query_api = self._client.query_api()
            result = query_api.query(self._query)
            
            # 解析查询结果
            data = {}
            for table in result:
                for record in table.records:
                    field = record.get_field()
                    value = record.get_value()
                    # 如果同一个字段有多个值，取最后一个
                    data[field] = value
            
            if data:
                self.dataReceived.emit(data)
                
        except Exception as e:
            error_msg = f"查询失败: {str(e)}"
            logger.error(error_msg)
            self.errorOccurred.emit(error_msg)
    
    @Slot()
    def disconnect(self):
        """断开连接"""
        # 安全地停止定时器（可能在程序退出时已被销毁）
        if self._timer:
            try:
                self._timer.stop()
            except RuntimeError:
                # QTimer 已被销毁，忽略错误
                pass
            self._timer = None
        
        if self._client:
            try:
                self._client.close()
            except Exception:
                # 客户端可能已被关闭，忽略错误
                pass
            self._client = None
        
        self._is_connected = False
    
    def __del__(self):
        """析构函数：安全地清理资源"""
        try:
            self.disconnect()
        except Exception:
            # 程序退出时 Qt 对象可能已被销毁，忽略所有异常
            pass