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