feat(HEL-382): 搭建 datahub 底座和盘后正式数据链路
新增独立 xiaobai-datahub 服务(SQLite WAL、Tushare 盘后发布、/v1 契约和管理后台),不改现站页面与数据链路。 Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: multica-agent <github@multica.ai>
This commit is contained in:
co-authored by
Cursor
multica-agent
parent
c2ebc0ab91
commit
3498dd7a4b
@@ -0,0 +1,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from datahub.adapters.tushare import TushareAdapter
|
||||
from datahub.crypto import SecretVault
|
||||
from datahub.db import HubDB
|
||||
from datahub.pipeline import Pipeline, QualityError
|
||||
from datahub.settings import Settings
|
||||
from tests.fixtures import TRADE_DATE, fake_transport
|
||||
|
||||
|
||||
def make_pipeline(before_commit=None) -> tuple[Pipeline, HubDB]:
|
||||
tmp = tempfile.TemporaryDirectory()
|
||||
db = HubDB(Path(tmp.name) / "hub.db")
|
||||
adapter = TushareAdapter("test-token", transport=fake_transport)
|
||||
settings = Settings(
|
||||
encryption_key=SecretVault.generate_key(),
|
||||
api_token="t" * 32,
|
||||
admin_password="admin-pass",
|
||||
tushare_token="test-token",
|
||||
db_path=db.path,
|
||||
quality={"daily_row_ratio": 0.98, "null_rate_max": 0.01, "max_publish_attempts": 3, "publication_generations": 3},
|
||||
scheduler_enabled=False,
|
||||
)
|
||||
pipe = Pipeline(db, adapter, settings, before_commit=before_commit)
|
||||
pipe._tmp = tmp # keep alive
|
||||
return pipe, db
|
||||
|
||||
|
||||
class PipelineTests(unittest.TestCase):
|
||||
def test_reference_and_daily_publish(self) -> None:
|
||||
pipe, db = make_pipeline()
|
||||
ref = pipe.ingest_reference(TRADE_DATE)
|
||||
self.assertEqual(ref["stocks"], 2)
|
||||
result = pipe.run_dataset("daily", TRADE_DATE)
|
||||
self.assertEqual(result["state"], "published")
|
||||
self.assertEqual(result["rows"], 2)
|
||||
pub = db.fetchone("SELECT * FROM publications WHERE dataset='daily' AND trade_date=?", (TRADE_DATE,))
|
||||
self.assertEqual(pub["active_batch"], result["batch_id"])
|
||||
rows = db.fetchall("SELECT * FROM eod_bars WHERE batch_id=?", (result["batch_id"],))
|
||||
self.assertEqual(len(rows), 2)
|
||||
self.assertEqual(rows[0]["amount"] if rows[0]["ts_code"] == "600000.SH" else rows[1]["amount"], 2_000_000.0)
|
||||
|
||||
def test_atomic_publish_abort_leaves_no_half_batch(self) -> None:
|
||||
pipe, db = make_pipeline()
|
||||
pipe.ingest_reference(TRADE_DATE)
|
||||
first = pipe.run_dataset("daily", TRADE_DATE)
|
||||
boom = {"n": 0}
|
||||
|
||||
def explode() -> None:
|
||||
boom["n"] += 1
|
||||
raise RuntimeError("killed")
|
||||
|
||||
pipe.before_commit = explode
|
||||
with self.assertRaises(RuntimeError):
|
||||
pipe.run_dataset("daily", TRADE_DATE)
|
||||
pub = db.fetchone("SELECT * FROM publications WHERE dataset='daily' AND trade_date=?", (TRADE_DATE,))
|
||||
self.assertEqual(pub["active_batch"], first["batch_id"])
|
||||
visible = db.fetchall(
|
||||
"SELECT DISTINCT batch_id FROM eod_bars WHERE trade_date=? AND batch_id=?",
|
||||
(TRADE_DATE, pub["active_batch"]),
|
||||
)
|
||||
self.assertEqual(len(visible), 1)
|
||||
|
||||
def test_rollback_switches_active_batch(self) -> None:
|
||||
pipe, _db = make_pipeline()
|
||||
pipe.ingest_reference(TRADE_DATE)
|
||||
first = pipe.run_dataset("daily", TRADE_DATE)
|
||||
second = pipe.run_dataset("daily", TRADE_DATE)
|
||||
self.assertNotEqual(first["batch_id"], second["batch_id"])
|
||||
rolled = pipe.rollback("daily", TRADE_DATE, actor="test")
|
||||
self.assertEqual(rolled["active_batch"], first["batch_id"])
|
||||
from datahub.serving import V1API
|
||||
|
||||
api = V1API(pipe.db, pipe, pipe.settings)
|
||||
payload = api.handle("/v1/bars/daily", {"date": [TRADE_DATE], "code": ["600000.SH"]})
|
||||
self.assertEqual(payload["meta"]["batch_id"], first["batch_id"])
|
||||
|
||||
def test_row_ratio_gate_rejects_short_batch(self) -> None:
|
||||
pipe, _db = make_pipeline()
|
||||
pipe.ingest_reference(TRADE_DATE)
|
||||
original = fake_transport
|
||||
|
||||
def short(api_name, params, fields):
|
||||
rows = original(api_name, params, fields)
|
||||
if api_name == "daily":
|
||||
return rows[:1]
|
||||
return rows
|
||||
|
||||
pipe.adapter._transport = short
|
||||
with self.assertRaises(QualityError) as ctx:
|
||||
pipe.run_dataset("daily", TRADE_DATE)
|
||||
self.assertTrue(ctx.exception.report["hard_fail"])
|
||||
pub = pipe.db.fetchone("SELECT * FROM publications WHERE dataset='daily'")
|
||||
self.assertIsNone(pub)
|
||||
|
||||
def test_wal_mode(self) -> None:
|
||||
pipe, db = make_pipeline()
|
||||
with db.connect() as connection:
|
||||
mode = connection.execute("PRAGMA journal_mode").fetchone()[0]
|
||||
self.assertEqual(str(mode).lower(), "wal")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user