feat: add TDengine connector + schema modules with tests

This commit is contained in:
张宗平
2026-06-11 15:38:59 +08:00
parent 0788daf2af
commit 7d1c03b98d
3 changed files with 276 additions and 0 deletions
+117
View File
@@ -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:
"""执行非查询 SQLINSERT / 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)