feat: 完成 issue #4 数据总线 + 时序库 模板化封装

This commit is contained in:
2026-08-04 16:58:06 +08:00
parent 1fb1d278d5
commit 4c3a9fdbe6
10 changed files with 789 additions and 3 deletions
+43
View File
@@ -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
```
+43 -1
View File
@@ -5,5 +5,47 @@
复用 Kafka + TDengine + PostgreSQL + MinIO,全部改为**按模板配置** 复用 Kafka + TDengine + PostgreSQL + MinIO,全部改为**按模板配置**
topic / 时序库表 / 关系表 / 对象桶,换行业只改模板资产(YAML + 点位字典), 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",
]
+26
View File
@@ -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)])
+232
View File
@@ -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()
+2 -2
View File
@@ -25,7 +25,7 @@ TABLE_COLUMNS: dict = {
"id BIGSERIAL PRIMARY KEY", "id BIGSERIAL PRIMARY KEY",
"template VARCHAR(64) NOT NULL UNIQUE", "template VARCHAR(64) NOT NULL UNIQUE",
"version VARCHAR(32) NOT NULL DEFAULT '0.1.0'", "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'", "status VARCHAR(16) NOT NULL DEFAULT 'draft'",
"created_at TIMESTAMPTZ NOT NULL DEFAULT now()", "created_at TIMESTAMPTZ NOT NULL DEFAULT now()",
], ],
@@ -34,7 +34,7 @@ TABLE_COLUMNS: dict = {
"model_id VARCHAR(128) NOT NULL UNIQUE", "model_id VARCHAR(128) NOT NULL UNIQUE",
"template VARCHAR(64) REFERENCES {schema}.templates(template) ON DELETE CASCADE", "template VARCHAR(64) REFERENCES {schema}.templates(template) ON DELETE CASCADE",
"algorithm VARCHAR(32) NOT NULL DEFAULT 'xgboost'", "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'", "version VARCHAR(32) NOT NULL DEFAULT '0.1.0'",
"status VARCHAR(16) NOT NULL DEFAULT 'staging'", "status VARCHAR(16) NOT NULL DEFAULT 'staging'",
"created_at TIMESTAMPTZ NOT NULL DEFAULT now()", "created_at TIMESTAMPTZ NOT NULL DEFAULT now()",
+16
View File
@@ -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
+223
View File
@@ -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()
@@ -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()
@@ -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()
+89
View File
@@ -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()