feat: add watch subcommand (TMQ + polling fallback)

This commit is contained in:
张宗平
2026-06-11 16:12:38 +08:00
parent cdfd77bdcb
commit 7e6915ade6
4 changed files with 198 additions and 0 deletions
+3
View File
@@ -22,6 +22,9 @@ python -m ts_anomaly_td forecast-anomaly --stable finance_001 --output render/
# 端到端
python -m ts_anomaly_td e2e --data-dir data --output-dir render --log-dir logs/
# 准实时监控
python -m ts_anomaly_td watch --stable finance_001 --window-sec 30
```
## E2E 步骤
+55
View File
@@ -0,0 +1,55 @@
"""tmq_watch 模块测试(mock TMQ consumer)。"""
from unittest.mock import MagicMock, patch
import pytest
from ts_anomaly_td.tmq_watch import WatchRunner
def make_mock_conn():
conn = MagicMock()
conn.execute.return_value = [(1000, 1.0), (2000, 2.0)]
return conn
def test_watch_runner_polling_fallback(monkeypatch):
"""TMQ 不可用时自动降级为轮询。"""
conn = make_mock_conn()
runner = WatchRunner(conn, "test", window_sec=0.1)
# 强制 TMQ 不可用
runner._run_tmq = MagicMock(side_effect=ImportError("no TMQ"))
# 跑一个轮询周期就停
poll_count = [0]
def mock_poll():
poll_count[0] += 1
runner._mode = "polling"
runner.stop()
runner._run_polling = mock_poll
runner.run()
assert runner._mode == "polling"
assert poll_count[0] == 1
def test_watch_runner_stop_on_signal():
"""收到 stop 后退出循环。"""
conn = make_mock_conn()
runner = WatchRunner(conn, "test", window_sec=0.1)
runner.stop()
assert runner._stop is True
def test_watch_runner_window_check():
"""_check_window 调用 detection。"""
conn = make_mock_conn()
runner = WatchRunner(conn, "test")
with patch("ts_anomaly_td.detection.detect_all_algos") as mock_det:
mock_det.return_value = {
"ksigma": {"windows": [(1000, 2000)], "error": None},
}
runner._check_window([(1000, 1.0)])
mock_det.assert_called_once_with(conn, "test")
+24
View File
@@ -79,6 +79,22 @@ def cmd_e2e(args):
return run_e2e(args)
def cmd_watch(args):
"""watch 子命令:TMQ 准实时监控。"""
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
from ts_anomaly_td.tmq_watch import WatchRunner
conn = TDConnection(args.url)
try:
db = args.db or "ts_anomaly"
conn.execute_no_result(f"USE {db}")
runner = WatchRunner(conn, args.stable, window_sec=args.window_sec)
runner.run()
finally:
conn.close()
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog="ts-anomaly-td",
@@ -116,6 +132,13 @@ def build_parser() -> argparse.ArgumentParser:
p_e2e.add_argument("--log-dir", default="logs", help="JSON 日志目录 (default: logs)")
p_e2e.add_argument("--url", default="ws://root:taosdata@localhost:6041", help="TDengine WebSocket URL")
# watch (Iter-2)
p_watch = sub.add_parser("watch", help="TMQ 准实时异常监控")
p_watch.add_argument("--stable", required=True, help="supertable 名称")
p_watch.add_argument("--db", default="ts_anomaly", help="数据库名")
p_watch.add_argument("--window-sec", type=int, default=30, help="检测间隔秒数 (default: 30)")
p_watch.add_argument("--url", default="ws://root:taosdata@localhost:6041", help="TDengine WebSocket URL")
return parser
@@ -128,6 +151,7 @@ def main():
"detect-batch": cmd_detect_batch,
"forecast-anomaly": cmd_forecast_anomaly,
"e2e": cmd_e2e,
"watch": cmd_watch,
}
fn = dispatch.get(args.command)
+116
View File
@@ -0,0 +1,116 @@
"""TMQ 准实时订阅模块:订阅 supertable topic + 周期 ANOMALY_WINDOW。"""
import logging
import signal
import time
logger = logging.getLogger(__name__)
class WatchRunner:
"""TMQ watch 运行器(带轮询降级)。
Args:
conn: TDConnection 实例
stable: supertable 名称
window_sec: ANOMALY_WINDOW 触发间隔(秒),默认 30
window_size: 滑动窗口大小(条),默认 1000
"""
def __init__(self, conn, stable: str, window_sec: int = 30, window_size: int = 1000):
self.conn = conn
self.stable = stable
self.window_sec = window_sec
self.window_size = window_size
self._stop = False
self._mode = "unknown"
def stop(self):
"""优雅停止。"""
self._stop = True
def run(self):
"""主循环。先尝试 TMQ,失败则降级为轮询。"""
signal.signal(signal.SIGINT, self._handle_signal)
signal.signal(signal.SIGTERM, self._handle_signal)
try:
self._run_tmq()
except Exception as e:
logger.warning("TMQ unavailable (%s), falling back to polling", e)
print("watch: using polling mode (TMQ unavailable)")
self._run_polling()
def _handle_signal(self, signum, frame):
logger.info("received signal %d, shutting down", signum)
self.stop()
def _run_tmq(self):
"""TMQ 模式:订阅 supertable topic → 滑动窗口 → 周期 ANOMALY_WINDOW。"""
self._mode = "tmq"
import taosws
consumer = taosws.Consumer()
consumer.subscribe([f"ds_{self.stable}"])
buffer = []
last_check = time.monotonic()
while not self._stop:
try:
records = consumer.consume(timeout=2.0)
for record in records:
buffer.append((record.ts, record.value))
if len(buffer) > self.window_size:
buffer = buffer[-self.window_size:]
now = time.monotonic()
if now - last_check >= self.window_sec and buffer:
self._check_window(buffer)
last_check = now
except Exception as e:
logger.error("TMQ consume error: %s", e)
time.sleep(1)
consumer.unsubscribe()
consumer.close()
logger.info("watch stopped (TMQ mode)")
def _run_polling(self):
"""轮询模式:周期 SELECT + ANOMALY_WINDOW。"""
self._mode = "polling"
while not self._stop:
try:
sql = (
f"SELECT ts, value FROM s_{self.stable} "
f"ORDER BY ts DESC LIMIT {self.window_size}"
)
rows = self.conn.execute(sql)
buffer = [(int(r[0]), float(r[1])) for r in rows]
if buffer:
self._check_window(buffer)
except Exception as e:
logger.error("polling error: %s", e)
time.sleep(self.window_sec)
logger.info("watch stopped (polling mode)")
def _check_window(self, buffer: list):
"""对缓冲数据跑 ANOMALY_WINDOW + ANSI 输出。"""
try:
from ts_anomaly_td.detection import detect_all_algos
results = detect_all_algos(self.conn, self.stable)
for algo, r in results.items():
if r["error"]:
continue
for wstart, wend in r["windows"]:
print(
f"\033[33m[ANOMALY]\033[0m "
f"\033[36m{algo}\033[0m "
f"{wstart}-{wend}"
)
except Exception as e:
logger.error("check_window failed: %s", e)