张宗平
118 lines
3.8 KiB
Python
118 lines
3.8 KiB
Python
"""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)
|