diff --git a/tests/test_connector.py b/tests/test_connector.py new file mode 100644 index 0000000..f0f8c39 --- /dev/null +++ b/tests/test_connector.py @@ -0,0 +1,99 @@ +"""connector 模块测试(mock taosws.connect)。""" + +from unittest.mock import patch, MagicMock +import pytest +from ts_anomaly_td.connector import TDConnection, ConnectionError, _MAX_RETRY_DURATION + + +@pytest.fixture +def mock_connect(): + """mock taosws.connect 返回假连接。""" + with patch("ts_anomaly_td.connector.taosws.connect") as m: + mock_conn = MagicMock() + mock_conn.execute.return_value = [[("3.3.6.0",)]] + m.return_value = mock_conn + yield m + + +def test_connect_success(mock_connect): + """首次连接成功。""" + conn = TDConnection() + assert conn.is_healthy() + mock_connect.assert_called_once() + + +def test_execute_query(mock_connect): + """execute 返回行列表。""" + mock_conn = mock_connect.return_value + # taosws 返回可迭代的行对象;每行 row 是 (1, 'a') / (2, 'b') + mock_conn.execute.return_value = [(1, "a"), (2, "b")] + + conn = TDConnection() + rows = conn.execute("SELECT 1") + assert rows == [(1, "a"), (2, "b")] + + +def test_execute_no_result(mock_connect): + """execute_no_result 不返回值。""" + conn = TDConnection() + result = conn.execute_no_result("CREATE DATABASE test") + assert result is None + + +def test_retry_logic(mock_connect): + """前 2 次失败、第 3 次成功。""" + import ts_anomaly_td.connector as mod + + # 缩短退避以加速测试 + mod._MAX_BACKOFF = 0.01 + try: + mock_connect.side_effect = [ + Exception("connection refused"), + Exception("connection refused"), + MagicMock(execute=MagicMock(return_value=[[("3.3.6.0",)]])), + ] + + conn = TDConnection() + assert conn.is_healthy() + assert mock_connect.call_count == 3 + finally: + mod._MAX_BACKOFF = 300 + + +def test_retry_timeout(mock_connect): + """持续失败超过 30 分钟应抛出 ConnectionError。""" + import ts_anomaly_td.connector as mod + + mod._MAX_RETRY_DURATION = 0.1 + mod._MAX_BACKOFF = 0.01 + try: + mock_connect.side_effect = Exception("always fails") + + with pytest.raises(ConnectionError, match="30"): + conn = TDConnection() + # 触发懒加载连接:访问 .conn 属性(非 is_healthy,后者会吞掉异常) + _ = conn.conn + finally: + mod._MAX_RETRY_DURATION = 1800 + mod._MAX_BACKOFF = 300 + + +def test_close(mock_connect): + """close 关闭底层连接。""" + conn = TDConnection() + # 触发懒加载连接 + conn.is_healthy() + conn.close() + mock_connect.return_value.close.assert_called_once() + + +def test_batch_insert(mock_connect): + """batch_insert 拼接 INSERT SQL。""" + conn = TDConnection() + rows = [(1000, 1.0, 0), (2000, 2.0, 1)] + conn.batch_insert("test", rows) + + mock_conn = mock_connect.return_value + # 验证 execute 被调用 + calls = [str(c[0][0]) for c in mock_conn.execute.call_args_list] + assert any("INSERT INTO s_test VALUES" in c for c in calls) diff --git a/ts_anomaly_td/connector.py b/ts_anomaly_td/connector.py new file mode 100644 index 0000000..3ece99a --- /dev/null +++ b/ts_anomaly_td/connector.py @@ -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) diff --git a/ts_anomaly_td/schema.py b/ts_anomaly_td/schema.py new file mode 100644 index 0000000..2e1ffd5 --- /dev/null +++ b/ts_anomaly_td/schema.py @@ -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)