feat: add TDengine connector + schema modules with tests
This commit is contained in:
张宗平
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user