feat: 完成 issue #4 数据总线 + 时序库 模板化封装
This commit is contained in:
@@ -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
|
||||||
|
```
|
||||||
@@ -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",
|
||||||
|
]
|
||||||
|
|||||||
@@ -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)])
|
||||||
@@ -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()
|
||||||
@@ -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()",
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user