From 7e6915ade6ba781412edd1a2a73e38318dbd7bd3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BC=A0=E5=AE=97=E5=B9=B3?= Date: Thu, 11 Jun 2026 16:12:38 +0800 Subject: [PATCH] feat: add watch subcommand (TMQ + polling fallback) --- README.md | 3 + tests/test_tmq_watch.py | 55 ++++++++++++++++++ ts_anomaly_td/cli.py | 24 ++++++++ ts_anomaly_td/tmq_watch.py | 116 +++++++++++++++++++++++++++++++++++++ 4 files changed, 198 insertions(+) create mode 100644 tests/test_tmq_watch.py create mode 100644 ts_anomaly_td/tmq_watch.py diff --git a/README.md b/README.md index 861caea..f376fc8 100644 --- a/README.md +++ b/README.md @@ -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 步骤 diff --git a/tests/test_tmq_watch.py b/tests/test_tmq_watch.py new file mode 100644 index 0000000..8496f9a --- /dev/null +++ b/tests/test_tmq_watch.py @@ -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") diff --git a/ts_anomaly_td/cli.py b/ts_anomaly_td/cli.py index 9848bbb..efc858b 100644 --- a/ts_anomaly_td/cli.py +++ b/ts_anomaly_td/cli.py @@ -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) diff --git a/ts_anomaly_td/tmq_watch.py b/ts_anomaly_td/tmq_watch.py new file mode 100644 index 0000000..1d39044 --- /dev/null +++ b/ts_anomaly_td/tmq_watch.py @@ -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)