diff --git a/core/data-bus/README.md b/core/data-bus/README.md new file mode 100644 index 0000000..bf85a52 --- /dev/null +++ b/core/data-bus/README.md @@ -0,0 +1,43 @@ +# Data Bus —— 数据总线 + 时序库模板化封装 + +对应 PRD 5.2「② 数据总线 + 时序库」与 Issue #4(EPIC)。 +复用 **Kafka + TDengine + PostgreSQL + MinIO**,全部改为**按模板配置** +topic / 时序库表 / 关系表 / 对象桶:换行业只改模板资产(模板 YAML + 点位字典 CSV), +内核代码零改动。 + +## 模块 + +| 文件 | 职责 | +|------|------| +| `templating.py` | 模板命名推导:Kafka topic + 分区策略、TDengine 超级表/子表、PostgreSQL schema/表、MinIO 桶/对象键 | +| `tdengine_schema.py` | 时序库 schema 自动生成(超级表 + 每测点子表)+ 批量 INSERT SQL | +| `postgres_schema.py` | 关系库 schema(模板 / 模型 / 用户 / 权限)+ 角色授权语句 | +| `batch_writer.py` | 批量写入缓冲:批量聚合(默认 5000 条/0.1s,对齐 PRD 5.2「5k/100ms」基线)、幂等去重、失败重试 —— **数据不丢不重** | + +## 使用示例 + +```python +from data_bus.batch_writer import BatchWriter, TdengineSink +from data_bus.templating import TemplateNaming + +naming = TemplateNaming(template="ti-cl4") # 换行业只改模板名 +sink = TdengineSink(naming, executor=run_sql) # executor 注入 TDengine 连接适配器 +writer = BatchWriter(sink, batch_size=5000, flush_interval=0.1) + +writer.push_many(rows) # 批量入队(重复键自动过滤) +writer.flush() # 主动刷盘;失败保留缓冲重发 +``` + +## 验收口径(Issue #4) + +- **端到端写入**:`push → flush → sink` 全链路(MemorySink 本地联调 / TdengineSink 落库)。 +- **批量写入**:batch_size / flush_interval 触发批量 flush(默认 5000 条 / 0.1s)。 +- **数据不丢不重**:写成功才清缓冲(不丢);(device_id, point_id, ts) 幂等去重 + (不重,含失败重发的部分写入场景,sink 侧二次兜底)。 + +## 测试 + +```bash +cd core/data-bus +python -m unittest discover -s tests -v +``` diff --git a/core/data-bus/__init__.py b/core/data-bus/__init__.py index 96326fc..c79856e 100644 --- a/core/data-bus/__init__.py +++ b/core/data-bus/__init__.py @@ -5,5 +5,47 @@ 复用 Kafka + TDengine + PostgreSQL + MinIO,全部改为**按模板配置** topic / 时序库表 / 关系表 / 对象桶,换行业只改模板资产(YAML + 点位字典), 内核代码零改动;批量写入缓冲 + 幂等去重,保障「数据不丢不重」。 + +模块: +- templating 模板命名推导(Kafka topic/分区、TDengine 表、PG schema、MinIO 桶); +- tdengine_schema 时序超级表/子表 DDL + 批量 INSERT SQL 生成; +- postgres_schema 关系库 schema(模板/模型/用户/权限)+ 授权语句; +- batch_writer 批量写入缓冲(StoreSink 抽象 / MemorySink / TdengineSink), + 批量聚合 + 幂等去重 + 失败重试(不丢不重)。 + +测试:`python -m unittest discover -s tests -v`(在 core/data-bus 目录下执行)。 """ -__version__ = "0.1.0" +__version__ = "0.2.0" + +from .batch_writer import ( + DEDUP_KEYS, + BatchWriter, + MemorySink, + StoreSink, + TdengineSink, +) +from .postgres_schema import generate_grant_ddl, generate_schema_ddl +from .tdengine_schema import ( + PointSpec, + build_batch_insert, + generate_subtable_ddls, + generate_supertable_ddl, +) +from .templating import TemplateNaming, sanitize, sanitize_sql + +__all__ = [ + "TemplateNaming", + "sanitize", + "sanitize_sql", + "PointSpec", + "generate_supertable_ddl", + "generate_subtable_ddls", + "build_batch_insert", + "generate_schema_ddl", + "generate_grant_ddl", + "StoreSink", + "MemorySink", + "TdengineSink", + "BatchWriter", + "DEDUP_KEYS", +] diff --git a/core/data-bus/_sanity_check.py b/core/data-bus/_sanity_check.py new file mode 100644 index 0000000..dd516c4 --- /dev/null +++ b/core/data-bus/_sanity_check.py @@ -0,0 +1,26 @@ +# -*- coding: utf-8 -*- +"""Sanity check: data_bus package imports and postgres DDL escaping fix.""" +import importlib.machinery +import importlib.util +import os +import sys + +d = os.getcwd() + +# 以包身份加载 data-bus(执行 __init__.py,相对导入可用) +spec = importlib.util.spec_from_file_location( + "data_bus", os.path.join(d, "__init__.py"), submodule_search_locations=[d] +) +pkg = importlib.util.module_from_spec(spec) +sys.modules["data_bus"] = pkg +spec.loader.exec_module(pkg) + +from data_bus import BatchWriter, MemorySink, TdengineSink, TemplateNaming, __version__ # noqa: E402 +from data_bus.postgres_schema import generate_schema_ddl # noqa: E402 + +ddl = generate_schema_ddl(TemplateNaming(template="ti-cl4")) +assert "DEFAULT '{}'::jsonb" in ddl, "jsonb default escaped wrongly" +assert "REFERENCES tpl_ti_cl4.templates(template)" in ddl +print("version:", __version__) +print("ddl ok; lines:", len(ddl.splitlines())) +print("exports ok:", [c.__name__ for c in (BatchWriter, MemorySink, TdengineSink)]) diff --git a/core/data-bus/batch_writer.py b/core/data-bus/batch_writer.py new file mode 100644 index 0000000..c9a42fd --- /dev/null +++ b/core/data-bus/batch_writer.py @@ -0,0 +1,232 @@ +# -*- coding: utf-8 -*- +"""批量写入缓冲 + 幂等去重 —— 「数据不丢不重」的实现保障(PRD 5.2 / Issue #4)。 + +设计: +- push() 单条/批量入队:按 (device_id, point_id, ts) 幂等去重(**不重**); +- 触发 flush:缓冲达到 batch_size 或距上次 flush 超过 flush_interval + (**批量写入**,默认 5000 条 / 0.1s,对齐 PRD 5.2「5k 条/100ms」基线); +- at-least-once:sink 写成功才清空缓冲;失败整批保留、下次 flush 重发(**不丢**); +- sink 侧按去重键幂等(MemorySink / TdengineSink 均已实现):失败重发时 + 即使发生部分写入也不会重复落库,以 sink 返回的成功条数为准。 + +StoreSink 抽象: +- MemorySink —— 内存落库(本地联调 / 端到端测试); +- TdengineSink —— 把样本批转为按子表分组的批量 INSERT,交客户端执行 + (不依赖 taospy;executor 由调用方注入连接适配器或联调函数)。 +""" +from __future__ import annotations + +import threading +import time +from abc import ABC, abstractmethod +from typing import Callable, Dict, List, Optional, Tuple + +from .tdengine_schema import build_batch_insert + +# 幂等去重键:同一设备-测点-时间戳视为同一条样本(重复推送只落一次) +DEDUP_KEYS: Tuple[str, ...] = ("device_id", "point_id", "ts") + + +def _dedup_key(row: dict, keys: Tuple[str, ...]) -> tuple: + """从样本行提取去重键(缺失字段统一为字符串 "None",保证键可哈希稳定)。""" + return tuple(str(row.get(k)) for k in keys) + + +class StoreSink(ABC): + """存储端抽象:批量写入的落库实现。""" + + @abstractmethod + def write(self, rows: List[dict]) -> int: + """写入一批样本。 + + Args: + rows: 样本列表(含 device_id / point_id / value / ts / quality)。 + + Returns: + 成功写入的条数。 + + Raises: + 写入失败抛异常 —— 调用方(BatchWriter)保留缓冲待重发(不丢)。 + """ + + def close(self) -> None: + """释放资源(可选覆写;幂等,可多次调用)。""" + + +class MemorySink(StoreSink): + """内存落库:按去重键幂等、行保序(本地联调 / 端到端测试用)。""" + + def __init__(self, dedup_keys: Tuple[str, ...] = DEDUP_KEYS): + self.dedup_keys = dedup_keys + self._rows: List[dict] = [] + self._seen: set = set() + self.write_count: int = 0 + self.closed: bool = False + + def write(self, rows: List[dict]) -> int: + n = 0 + for row in rows: + key = _dedup_key(row, self.dedup_keys) + if key in self._seen: + continue + self._seen.add(key) + self._rows.append(dict(row)) + n += 1 + self.write_count += n + return n + + def close(self) -> None: + self.closed = True + + @property + def rows(self) -> List[dict]: + """已落库样本(保序副本,供验收断言)。""" + return list(self._rows) + + +class TdengineSink(StoreSink): + """TDengine 批量落库:样本批 → 按子表分组的批量 INSERT,交 executor 执行。 + + 不依赖 taospy:executor(statements) 由调用方注入(连接适配器执行 SQL, + 或打印 SQL 的联调函数)。子表命名由模板命名器推导(点位字典驱动)。 + """ + + def __init__( + self, + naming: "TemplateNaming", + executor: Callable[[List[str]], None], + dedup_keys: Tuple[str, ...] = DEDUP_KEYS, + ): + self.naming = naming + self.executor = executor + self.dedup_keys = dedup_keys + self._seen: set = set() + self.write_count: int = 0 + + def write(self, rows: List[dict]) -> int: + fresh = [] + for row in rows: + key = _dedup_key(row, self.dedup_keys) + if key in self._seen: + continue + self._seen.add(key) + fresh.append(row) + if fresh: + # 批量 INSERT:每子表一条 VALUES 多值语句(对齐 PRD 5.2 批量写入) + self.executor(build_batch_insert(self.naming, fresh)) + self.write_count += len(fresh) + return len(fresh) + + +class BatchWriter: + """批量写入缓冲:批量聚合 + 幂等去重 + 失败重试(**不丢不重**)。 + + Args: + sink: StoreSink 落库实现(MemorySink / TdengineSink / 自定义)。 + batch_size: 缓冲达到该条数触发一次 flush(默认 5000,5k/100ms 基线)。 + flush_interval: 距上次 flush 超过该秒数,push 时自动 flush(默认 0.1s)。 + dedup_keys: 幂等去重键(默认 设备-测点-时间戳)。 + clock: 时钟注入(测试可传假时钟),默认 time.monotonic。 + """ + + def __init__( + self, + sink: StoreSink, + batch_size: int = 5000, + flush_interval: float = 0.1, + dedup_keys: Tuple[str, ...] = DEDUP_KEYS, + clock: Callable[[], float] = time.monotonic, + ): + self.sink = sink + self.batch_size = max(1, int(batch_size)) + self.flush_interval = max(0.0, float(flush_interval)) + self.dedup_keys = dedup_keys + self._clock = clock + self._buffer: List[dict] = [] + self._seen: set = set() # 已成功提交的去重键(跨批幂等窗口) + self._buffered_keys: set = set() # 当前缓冲内去重键(缓冲未提交前即拦截重复) + self._last_flush = self._clock() + self._stats: Dict[str, int] = { + "received": 0, + "written": 0, + "duplicates": 0, + "flush_count": 0, + "failed_flushes": 0, + } + self._lock = threading.Lock() + + # ------------------------------------------------------------------ + def push(self, row: dict) -> bool: + """入队一条样本。 + + 自动刷盘(达到 batch_size / 超 flush_interval)失败时**不中断入队**: + 失败批保留缓冲待重发(不丢),统计计入 failed_flushes,可经 stats() 观察。 + + Returns: + True=接受入队;False=与已入队/已提交样本重复(幂等丢弃,计入 duplicates)。 + """ + with self._lock: + key = _dedup_key(row, self.dedup_keys) + if key in self._seen or key in self._buffered_keys: + self._stats["duplicates"] += 1 + return False + self._buffer.append(dict(row)) + self._buffered_keys.add(key) + self._stats["received"] += 1 + if len(self._buffer) >= self.batch_size or ( + self._clock() - self._last_flush >= self.flush_interval + ): + try: + self._flush_locked() + except Exception: + pass # 失败保留缓冲,由下次 push / flush 重试 + return True + + def push_many(self, rows: List[dict]) -> int: + """批量入队,返回接受条数(重复自动过滤)。""" + accepted = 0 + for row in rows: + if self.push(row): + accepted += 1 + return accepted + + def flush(self) -> int: + """主动刷新:把缓冲整批写入 sink(失败抛异常并保留缓冲)。""" + with self._lock: + return self._flush_locked() + + def _flush_locked(self) -> int: + if not self._buffer: + return 0 + batch = list(self._buffer) + try: + n = self.sink.write(batch) + except Exception: + # 写入失败:缓冲整体保留,下次 flush 重发(不丢) + self._stats["failed_flushes"] += 1 + raise + # 写入成功:清缓冲;sink 侧已按去重键幂等,重发不会重复落库 + self._buffer.clear() + self._buffered_keys.clear() + for row in batch: + self._seen.add(_dedup_key(row, self.dedup_keys)) + self._stats["written"] += n + self._stats["flush_count"] += 1 + self._last_flush = self._clock() + return n + + def pending(self) -> int: + """当前缓冲中待写入条数。""" + with self._lock: + return len(self._buffer) + + def stats(self) -> dict: + """累计统计(received/written/duplicates/flush_count/failed_flushes)。""" + with self._lock: + return dict(self._stats) + + def close(self) -> None: + """收尾:刷新剩余缓冲并关闭 sink(写失败异常上抛,由调用方处理)。""" + with self._lock: + self._flush_locked() + self.sink.close() diff --git a/core/data-bus/postgres_schema.py b/core/data-bus/postgres_schema.py index b013904..5951a91 100644 --- a/core/data-bus/postgres_schema.py +++ b/core/data-bus/postgres_schema.py @@ -25,7 +25,7 @@ TABLE_COLUMNS: dict = { "id BIGSERIAL PRIMARY KEY", "template VARCHAR(64) NOT NULL UNIQUE", "version VARCHAR(32) NOT NULL DEFAULT '0.1.0'", - "config JSONB NOT NULL DEFAULT '{}'::jsonb", + "config JSONB NOT NULL DEFAULT '{{}}'::jsonb", # format() 转义 {} "status VARCHAR(16) NOT NULL DEFAULT 'draft'", "created_at TIMESTAMPTZ NOT NULL DEFAULT now()", ], @@ -34,7 +34,7 @@ TABLE_COLUMNS: dict = { "model_id VARCHAR(128) NOT NULL UNIQUE", "template VARCHAR(64) REFERENCES {schema}.templates(template) ON DELETE CASCADE", "algorithm VARCHAR(32) NOT NULL DEFAULT 'xgboost'", - "hyperparams JSONB NOT NULL DEFAULT '{}'::jsonb", + "hyperparams JSONB NOT NULL DEFAULT '{{}}'::jsonb", # format() 转义 {} "version VARCHAR(32) NOT NULL DEFAULT '0.1.0'", "status VARCHAR(16) NOT NULL DEFAULT 'staging'", "created_at TIMESTAMPTZ NOT NULL DEFAULT now()", diff --git a/core/data-bus/tests/_bootstrap.py b/core/data-bus/tests/_bootstrap.py new file mode 100644 index 0000000..96b5015 --- /dev/null +++ b/core/data-bus/tests/_bootstrap.py @@ -0,0 +1,16 @@ +# -*- coding: utf-8 -*- +"""测试引导:把 `core/data-bus` 以包名 `data_bus` 挂载到 sys.modules。 + +目录名 `data-bus` 含连字符,无法直接以包名 import;挂载后模块内相对导入 +(`from .templating import ...`)在 unittest 发现机制下可正常解析。 +""" +import os +import sys +import types + +DATA_BUS_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, DATA_BUS_DIR) +if "data_bus" not in sys.modules: + pkg = types.ModuleType("data_bus") + pkg.__path__ = [DATA_BUS_DIR] + sys.modules["data_bus"] = pkg diff --git a/core/data-bus/tests/test_batch_writer.py b/core/data-bus/tests/test_batch_writer.py new file mode 100644 index 0000000..0821ef3 --- /dev/null +++ b/core/data-bus/tests/test_batch_writer.py @@ -0,0 +1,223 @@ +# -*- coding: utf-8 -*- +"""批量写入缓冲(batch_writer)测试:批量聚合 / 幂等去重 / 失败重试 / 端到端。 + +覆盖 Issue #4 验收点:端到端写入、批量写入(PRD 5.2 基线参数)、数据不丢不重。 +""" +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import _bootstrap # noqa: F401 + +from data_bus.batch_writer import BatchWriter, MemorySink, TdengineSink +from data_bus.templating import TemplateNaming + + +class FakeClock: + """可控时钟:测试定时刷盘不依赖真实时间。""" + + def __init__(self, start=0.0): + self.now = start + + def __call__(self): + return self.now + + def advance(self, seconds): + self.now += seconds + + +class FlakySink(MemorySink): + """前 fail_times 次 write 抛异常,之后正常(模拟上游抖动)。""" + + def __init__(self, fail_times=1): + super().__init__() + self.fail_times = fail_times + self.failed_attempts = 0 + + def write(self, rows): + if self.fail_times > 0: + self.fail_times -= 1 + self.failed_attempts += 1 + raise RuntimeError("upstream unavailable") + return super().write(rows) + + +class PartialWriteSink(MemorySink): + """第一次 write 先落前 N 条再抛异常(模拟客户端部分写入后连接中断)。""" + + def __init__(self, partial=3): + super().__init__() + self.partial = partial + self.failed_attempts = 0 + + def write(self, rows): + if self.failed_attempts == 0: + self.failed_attempts += 1 + super().write(list(rows[: self.partial])) # 部分写入 + raise RuntimeError("connection lost after partial write") + return super().write(rows) + + +def make_row(i=0, point_id=None, ts=None): + """构造一条样本(对齐 edge-gateway 行格式 + batch 标签)。""" + return { + "device_id": f"CLF-{i % 5 + 1:02d}", + "point_id": point_id or f"CLF-{i % 5 + 1:02d}.P{i:03d}", + "value": 20.0 + i, + "ts": float(ts) if ts is not None else 1754294400.0 + i, + "quality": 1, + "batch": f"b{int(i / 100)}", + } + + +class BatchFlushTest(unittest.TestCase): + def test_flush_by_size(self): + sink = MemorySink() + writer = BatchWriter(sink, batch_size=250, flush_interval=999) + for i in range(1000): + writer.push(make_row(i)) + self.assertEqual(writer.stats()["flush_count"], 4) + self.assertEqual(writer.stats()["written"], 1000) + self.assertEqual(len(sink.rows), 1000) + self.assertEqual(writer.pending(), 0) + + def test_push_many_accepted_count(self): + sink = MemorySink() + writer = BatchWriter(sink, batch_size=100, flush_interval=999) + accepted = writer.push_many([make_row(i) for i in range(50)]) + self.assertEqual(accepted, 50) + self.assertEqual(writer.pending(), 50) + + def test_close_flushes_pending(self): + sink = MemorySink() + writer = BatchWriter(sink, batch_size=1000, flush_interval=999) + writer.push_many([make_row(i) for i in range(7)]) + writer.close() + self.assertEqual(writer.stats()["written"], 7) + self.assertTrue(sink.closed) + + +class DedupTest(unittest.TestCase): + def test_duplicate_push_rejected(self): + sink = MemorySink() + writer = BatchWriter(sink, batch_size=1000, flush_interval=999) + row = make_row(0) + self.assertTrue(writer.push(row)) + self.assertFalse(writer.push(dict(row))) # 同 设备-测点-ts → 幂等丢弃 + writer.flush() + self.assertEqual(writer.stats()["duplicates"], 1) + self.assertEqual(len(sink.rows), 1) + + def test_dedup_across_batches(self): + sink = MemorySink() + writer = BatchWriter(sink, batch_size=50, flush_interval=999) + rows = [make_row(i) for i in range(100)] + writer.push_many(rows) + # 跨批重推同键:仍只落一次 + re_pushed = writer.push_many(rows[:20]) + self.assertEqual(re_pushed, 0) + writer.flush() + self.assertEqual(writer.stats()["written"], 100) + self.assertEqual(writer.stats()["duplicates"], 20) + self.assertEqual(len(sink.rows), 100) + + +class NoLossTest(unittest.TestCase): + def test_push_survives_failed_auto_flush(self): + # 自动刷盘失败不中断入队:失败批保留、下次 push 即自动重试(不丢) + sink = FlakySink(fail_times=1) + writer = BatchWriter(sink, batch_size=500, flush_interval=999) + writer.push_many([make_row(i) for i in range(1000)]) + # 第 500 条触发失败刷盘(批保留),第 501 条 push 时重试成功 501 条 + self.assertEqual(writer.stats()["failed_flushes"], 1) + self.assertEqual(writer.stats()["written"], 501) + self.assertEqual(writer.pending(), 499) # 其余继续在缓冲,未丢失 + writer.flush() # 显式刷盘清空 + self.assertEqual(writer.stats()["written"], 1000) + self.assertEqual(writer.pending(), 0) + self.assertEqual(len(sink.rows), 1000) # 不丢 + self.assertEqual(sink.failed_attempts, 1) + + def test_explicit_flush_retry(self): + # 显式 flush 失败时缓冲保留,再次 flush 重试成功 + sink = FlakySink(fail_times=1) + writer = BatchWriter(sink, batch_size=1000, flush_interval=999) + writer.push_many([make_row(i) for i in range(50)]) + with self.assertRaises(RuntimeError): + writer.flush() + self.assertEqual(writer.pending(), 50) # 失败批保留 + writer.flush() # 重试成功 + self.assertEqual(writer.stats()["written"], 50) + self.assertEqual(writer.pending(), 0) + self.assertEqual(len(sink.rows), 50) + + def test_no_dup_after_partial_retry(self): + # sink 内部幂等:失败批次部分已写,重发也不会重复落库(不重) + sink = PartialWriteSink(partial=3) + writer = BatchWriter(sink, batch_size=10, flush_interval=999) + writer.push_many([make_row(i) for i in range(20)]) + self.assertEqual(writer.stats()["failed_flushes"], 1) + writer.flush() + self.assertEqual(writer.pending(), 0) + self.assertEqual(len(sink.rows), 20) # 无重复 + keys = [(r["device_id"], r["point_id"], r["ts"]) for r in sink.rows] + self.assertEqual(len(keys), len(set(keys))) # 显式断言去重键唯一 + + +class TimeFlushTest(unittest.TestCase): + def test_flush_after_interval(self): + clock = FakeClock() + sink = MemorySink() + writer = BatchWriter(sink, batch_size=10000, flush_interval=0.1, clock=clock) + writer.push_many([make_row(i) for i in range(5)]) + self.assertEqual(writer.pending(), 5) # 未达批大小且未超间隔 + clock.advance(0.1) + writer.push(make_row(5)) # 超过 flush_interval → 自动刷盘 + self.assertEqual(writer.stats()["flush_count"], 1) + self.assertEqual(writer.stats()["written"], 6) + self.assertEqual(writer.pending(), 0) + + +class EndToEndTest(unittest.TestCase): + def test_memory_sink_round_trip(self): + sink = MemorySink() + writer = BatchWriter(sink, batch_size=5000, flush_interval=0.1) + rows = [make_row(i) for i in range(300)] + writer.push_many(rows) + writer.flush() + self.assertEqual(len(sink.rows), 300) + # 保序:首条与末条一致 + self.assertEqual(sink.rows[0]["device_id"], rows[0]["device_id"]) + self.assertEqual(sink.rows[-1]["point_id"], rows[-1]["point_id"]) + # 标签字段透传(设备-测点-批次-质量) + self.assertEqual({r["batch"] for r in sink.rows}, {"b0", "b1", "b2"}) + + def test_tdengine_sink_emits_batch_insert(self): + executed = [] + naming = TemplateNaming(template="ti-cl4") + sink = TdengineSink(naming, executor=lambda stmts: executed.extend(stmts)) + writer = BatchWriter(sink, batch_size=1000, flush_interval=999) + # 同测点 20 条 → 单子表多值批量 INSERT + rows = [ + { + "device_id": "CLF-01", + "point_id": "CLF-01.P000", + "value": float(i), + "ts": 1754294400.0 + i, + "quality": 1, + } + for i in range(20) + ] + writer.push_many(rows) + writer.flush() + self.assertEqual(sink.write_count, 20) + self.assertEqual(len(executed), 1) # 一张子表 → 一条多值语句 + stmt = executed[0] + self.assertTrue(stmt.startswith("INSERT INTO ti_cl4_pt_clf_01.p000 VALUES (")) + self.assertEqual(stmt.count("), ("), 19) # 20 个值元组 + self.assertTrue(stmt.endswith(");")) + + +if __name__ == "__main__": + unittest.main() diff --git a/core/data-bus/tests/test_postgres_schema.py b/core/data-bus/tests/test_postgres_schema.py new file mode 100644 index 0000000..06672de --- /dev/null +++ b/core/data-bus/tests/test_postgres_schema.py @@ -0,0 +1,39 @@ +# -*- coding: utf-8 -*- +"""PostgreSQL schema 生成(postgres_schema)单元测试:DDL + 授权语句。""" +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import _bootstrap # noqa: F401 + +from data_bus.postgres_schema import DEFAULT_TABLES, generate_grant_ddl, generate_schema_ddl +from data_bus.templating import TemplateNaming + +NAMING = TemplateNaming(template="ti-cl4") + + +class SchemaDDLTest(unittest.TestCase): + def test_schema_and_tables(self): + ddl = generate_schema_ddl(NAMING) + self.assertIn("CREATE SCHEMA IF NOT EXISTS tpl_ti_cl4;", ddl) + for table in DEFAULT_TABLES: + self.assertIn(f"CREATE TABLE IF NOT EXISTS tpl_ti_cl4.{table} (", ddl) + + def test_models_foreign_key_uses_schema(self): + ddl = generate_schema_ddl(NAMING) + self.assertIn("REFERENCES tpl_ti_cl4.templates(template)", ddl) + + +class GrantDDLTest(unittest.TestCase): + def test_roles_ro_rw(self): + grants = generate_grant_ddl(NAMING, roles=["databus"]) + text = "\n".join(grants) + self.assertIn("GRANT USAGE ON SCHEMA tpl_ti_cl4 TO databus_ro;", text) + self.assertIn("GRANT SELECT ON ALL TABLES IN SCHEMA tpl_ti_cl4 TO databus_ro;", text) + self.assertIn("GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA tpl_ti_cl4 TO databus_rw;", text) + self.assertIn("GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA tpl_ti_cl4 TO databus_rw;", text) + + +if __name__ == "__main__": + unittest.main() diff --git a/core/data-bus/tests/test_tdengine_schema.py b/core/data-bus/tests/test_tdengine_schema.py new file mode 100644 index 0000000..3ebd65e --- /dev/null +++ b/core/data-bus/tests/test_tdengine_schema.py @@ -0,0 +1,76 @@ +# -*- coding: utf-8 -*- +"""TDengine schema 生成(tdengine_schema)单元测试:DDL + 批量 INSERT。""" +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import _bootstrap # noqa: F401 + +from data_bus.tdengine_schema import ( + PointSpec, + build_batch_insert, + generate_subtable_ddls, + generate_supertable_ddl, +) +from data_bus.templating import TemplateNaming + +NAMING = TemplateNaming(template="ti-cl4") + + +def sample_point(device_id="CLF-01", point_id="CLF-01.TEMP", unit="℃", data_type="float"): + return PointSpec(device_id=device_id, point_id=point_id, unit=unit, data_type=data_type) + + +class SupertableDDLTest(unittest.TestCase): + def test_ddl_shape(self): + ddl = generate_supertable_ddl(NAMING) + self.assertIn("CREATE STABLE IF NOT EXISTS ti_cl4_points", ddl) + self.assertIn("ts TIMESTAMP", ddl) + self.assertIn("device_id NCHAR(64)", ddl) + + def test_retention_keep(self): + ddl = generate_supertable_ddl(NAMING, retention_days=90) + self.assertIn("KEEP(90)", ddl) + + +class SubtableDDLTest(unittest.TestCase): + def test_one_per_unique_point(self): + points = [sample_point(), sample_point(), sample_point(point_id="CLF-01.PRES")] + ddls = generate_subtable_ddls(NAMING, points) + self.assertEqual(len(ddls), 2) # 同点位只生成一张子表 + self.assertTrue(any("ti_cl4_pt_clf_01.temp" in d for d in ddls)) + self.assertTrue(any("USING ti_cl4_points" in d for d in ddls)) + self.assertTrue(any("TAGS ('CLF-01', '℃', 'float')" in d for d in ddls)) + + +class BatchInsertTest(unittest.TestCase): + def _row(self, point_id="CLF-01.TEMP", value=32.5, ts=1754294400.5, quality=1): + return { + "device_id": "CLF-01", + "point_id": point_id, + "value": value, + "ts": ts, + "quality": quality, + } + + def test_group_by_subtable_and_values(self): + rows = [self._row(), self._row(value=33.1, ts=1754294401.0), self._row(point_id="CLF-01.PRES")] + statements = build_batch_insert(NAMING, rows) + self.assertEqual(len(statements), 2) # 两个测点 → 两张子表各一条多值 INSERT + temp_stmt = [s for s in statements if "clf_01.temp" in s][0] + self.assertIn( + "INSERT INTO ti_cl4_pt_clf_01.temp VALUES (1754294400500, 32.5, 1), (1754294401000, 33.1, 1);", + temp_stmt, + ) + + def test_null_value(self): + statements = build_batch_insert(NAMING, [self._row(value=None)]) + self.assertIn("NULL", statements[0]) + + def test_empty_rows(self): + self.assertEqual(build_batch_insert(NAMING, []), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/core/data-bus/tests/test_templating.py b/core/data-bus/tests/test_templating.py new file mode 100644 index 0000000..6a8ae6b --- /dev/null +++ b/core/data-bus/tests/test_templating.py @@ -0,0 +1,89 @@ +# -*- coding: utf-8 -*- +"""模板化命名(templating)单元测试:Kafka / TDengine / PostgreSQL / MinIO。""" +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import _bootstrap # noqa: F401 + +from data_bus.templating import TemplateNaming, sanitize, sanitize_sql + + +class SanitizeTest(unittest.TestCase): + def test_lower_and_replace_unsafe(self): + # 小写 + 空格/中文等非安全字符替换为 `_`,随后 strip 首尾 `._` + self.assertEqual(sanitize("Ti-Cl4 模板"), "ti-cl4") + self.assertEqual(sanitize("A.B-c_d"), "a.b-c_d") + + def test_empty_fallback(self): + self.assertEqual(sanitize(""), "tpl") + self.assertEqual(sanitize("..."), "tpl") + + def test_sql_extra_hyphen(self): + # SQL 标识符额外把 `-` 换成 `_`,避免不加引号时报错 + self.assertEqual(sanitize_sql("ti-cl4"), "ti_cl4") + self.assertEqual(sanitize_sql("A.B-c"), "a.b_c") + + +class KafkaNamingTest(unittest.TestCase): + def setUp(self): + self.naming = TemplateNaming(template="ti-cl4", topic_prefix="ti-cl4", num_partitions=12) + + def test_topic_format(self): + self.assertEqual(self.naming.topic("CLF-01"), "ti-cl4.clf-01.points") + + def test_partition_deterministic_and_in_range(self): + p1 = self.naming.partition("CLF-01") + p2 = self.naming.partition("CLF-01") + self.assertEqual(p1, p2) + self.assertGreaterEqual(p1, 0) + self.assertLess(p1, 12) + + def test_partitions_mapping(self): + mapping = self.naming.partitions(["CLF-01", "CLF-02"]) + self.assertEqual(set(mapping), {"CLF-01", "CLF-02"}) + self.assertEqual(mapping["CLF-01"], self.naming.partition("CLF-01")) + + +class TdengineNamingTest(unittest.TestCase): + def setUp(self): + self.naming = TemplateNaming(template="ti-cl4") + + def test_stable_and_subtable(self): + self.assertEqual(self.naming.stable(), "ti_cl4_points") + # SQL 标识符保留 `.`(点位 ID 常用 `设备.测点` 形态) + self.assertEqual(self.naming.subtable("CLF-01.TEMP"), "ti_cl4_pt_clf_01.temp") + + +class PostgresNamingTest(unittest.TestCase): + def setUp(self): + self.naming = TemplateNaming(template="ti-cl4") + + def test_schema_and_table(self): + self.assertEqual(self.naming.pg_schema(), "tpl_ti_cl4") + self.assertEqual(self.naming.pg_table("models"), "tpl_ti_cl4.models") + + +class MinioNamingTest(unittest.TestCase): + def setUp(self): + self.naming = TemplateNaming(template="ti-cl4") + + def test_bucket(self): + self.assertEqual(self.naming.bucket(), "ti-cl4-artifacts") + + def test_snapshot_key(self): + self.assertEqual( + self.naming.snapshot_key("quality-forecast", "2026-08-04", 3), + "features/quality_forecast/2026-08-04/000003.jsonl", + ) + + def test_model_artifact_key(self): + self.assertEqual( + self.naming.model_artifact_key("quality-forecast", "v1.2"), + "models/quality_forecast/v1.2/model.bin", + ) + + +if __name__ == "__main__": + unittest.main()