feat: add TDengine connector + schema modules with tests

2026-06-11 15:38:59 +08:00
parent 0788daf2af
commit 7d1c03b98d
3 changed files with 276 additions and 0 deletions
@@ -0,0 +1,117 @@
+"""TDengine 连接模块：taos-ws-py WebSocket 封装 + 指数退避重试。"""
+
+import time
+import logging
+import taosws
+
+logger = logging.getLogger(__name__)
+
+_MAX_RETRY_DURATION = 1800      # 30 分钟
+_MAX_BACKOFF = 300              # 5 分钟
+_INSERT_BATCH_SIZE = 500
+
+
+class ConnectionError(Exception):
+    """TDengine 连接不可用时抛出。"""
+    pass
+
+
+class TDConnection:
+    """TDengine WebSocket 连接封装。
+
+    Args:
+        url: WebSocket 连接字符串，默认 ws://root:taosdata@localhost:6041
+    """
+
+    def __init__(self, url: str = "ws://root:taosdata@localhost:6041"):
+        self._url = url
+        self._conn: taosws.Connection | None = None
+
+    @property
+    def conn(self) -> taosws.Connection:
+        if self._conn is None:
+            self._conn = self._connect_with_retry()
+        return self._conn
+
+    def _connect_with_retry(self) -> taosws.Connection:
+        """指数退避重试连接 TDengine。
+
+        初始退避 1s，倍增因子 2，最大间隔 300s，最长持续 1800s。
+        每次重试前执行 SELECT SERVER_VERSION() 验证服务可用。
+
+        Raises:
+            ConnectionError: 30 分钟内仍不可用
+        """
+        backoff = 1.0
+        started_at = time.monotonic()
+        last_exc = None
+
+        while True:
+            elapsed = time.monotonic() - started_at
+            if elapsed > _MAX_RETRY_DURATION:
+                raise ConnectionError(
+                    f"TDengine 连接失败，已重试 {elapsed:.0f}s（超过 30 分钟上限）"
+                ) from last_exc
+
+            try:
+                conn = taosws.connect(self._url)
+                conn.execute("SELECT SERVER_VERSION()")
+                logger.info("TDengine 连接成功 (url=%s)", self._url)
+                return conn
+            except Exception as e:
+                last_exc = e
+                logger.warning(
+                    "TDengine 连接失败 (attempt, backoff=%.1fs, elapsed=%.0fs): %s",
+                    backoff, elapsed, e,
+                )
+                time.sleep(backoff)
+                backoff = min(backoff * 2, _MAX_BACKOFF)
+
+    def execute(self, sql: str) -> list[tuple]:
+        """执行查询 SQL，返回结果行列表。"""
+        logger.debug("execute: %s", sql[:120])
+        result = self.conn.execute(sql)
+        rows = []
+        if result is not None:
+            for row in result:
+                rows.append(tuple(row))
+        return rows
+
+    def execute_no_result(self, sql: str) -> None:
+        """执行非查询 SQL（INSERT / DDL）。"""
+        logger.debug("execute_no_result: %s", sql[:120])
+        self.conn.execute(sql)
+
+    def close(self):
+        """关闭连接。"""
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None
+
+    def is_healthy(self) -> bool:
+        """检查连接健康状态。"""
+        try:
+            self.execute("SELECT SERVER_VERSION()")
+            return True
+        except Exception:
+            return False
+
+    def batch_insert(self, stable: str, rows: list[tuple[int, float, int]]) -> None:
+        """批量 INSERT 数据到 subtable。
+
+        每 500 条一个 batch，单次 execute_no_result()。
+
+        Args:
+            stable: supertable 名称（不含 ds_ 前缀）
+            rows: [(ts_ms, value, label), ...]
+        """
+        total = len(rows)
+        for i in range(0, total, _INSERT_BATCH_SIZE):
+            batch = rows[i:i + _INSERT_BATCH_SIZE]
+            values_clauses = []
+            for ts, val, label in batch:
+                values_clauses.append(f"({ts}, {val}, {label})")
+            sql = f"INSERT INTO s_{stable} VALUES {' '.join(values_clauses)}"
+            self.execute_no_result(sql)
+            logger.debug("batch_insert: %d/%d rows", min(i + _INSERT_BATCH_SIZE, total), total)
+        logger.info("batch_insert complete: %d rows into s_%s", total, stable)
@@ -0,0 +1,60 @@
+"""TDengine DDL 模块：supertable / subtable 创建。"""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def create_database(conn, db_name: str = "ts_anomaly") -> None:
+    """创建数据库（不存在时创建）。"""
+    conn.execute_no_result(f"CREATE DATABASE IF NOT EXISTS {db_name} KEEP 36500")
+    logger.info("database %s ready", db_name)
+
+
+def use_database(conn, db_name: str = "ts_anomaly") -> None:
+    """切换到指定数据库。"""
+    conn.execute_no_result(f"USE {db_name}")
+
+
+def create_supertable(conn, stable: str) -> None:
+    """创建 supertable（不存在时创建）。
+
+    Args:
+        stable: 数据集名称，如 finance_001
+    """
+    sql = (
+        f"CREATE STABLE IF NOT EXISTS ds_{stable} "
+        f"(ts TIMESTAMP, value DOUBLE, label INT) "
+        f"TAGS (series_id INT)"
+    )
+    conn.execute_no_result(sql)
+    logger.info("supertable ds_%s ready", stable)
+
+
+def create_subtable(conn, stable: str, series_id: int = 1) -> None:
+    """创建 subtable（不存在时创建）。
+
+    Args:
+        stable: 数据集名称
+        series_id: 时序 ID（单变量数据集为 1）
+    """
+    sql = (
+        f"CREATE TABLE IF NOT EXISTS s_{stable} "
+        f"USING ds_{stable} TAGS ({series_id})"
+    )
+    conn.execute_no_result(sql)
+    logger.info("subtable s_%s ready (series_id=%d)", stable, series_id)
+
+
+def setup_schema(conn, stable: str, db_name: str = "ts_anomaly") -> None:
+    """一键建库建表。
+
+    Args:
+        conn: TDConnection 实例
+        stable: 数据集名称
+        db_name: 数据库名称
+    """
+    create_database(conn, db_name)
+    use_database(conn, db_name)
+    create_supertable(conn, stable)
+    create_subtable(conn, stable)