Initial commit: InternalAuditInterprise

backend/app/datahub/bootstrap.py

@@ -0,0 +1,58 @@
+"""数据中台 schema 初始化。
+
+MVP 阶段以 SQLAlchemy metadata 建表（后续可迁移到 Alembic）。
+扩展按可用性可选启用：
+- btree_gist / vector：若可用则创建。
+- timescaledb：若可用则把 metric_event 转为超表；不可用则保持普通表（带时间索引）。
+"""
+
+from __future__ import annotations
+
+from sqlalchemy import text
+from sqlalchemy.engine import Engine
+
+from app.datahub import models  # noqa: F401  确保模型注册到 metadata
+from app.db import Base, get_engine
+
+
+def _extension_available(engine: Engine, name: str) -> bool:
+    with engine.connect() as conn:
+        row = conn.execute(
+            text("SELECT 1 FROM pg_available_extensions WHERE name = :n"), {"n": name}
+        ).first()
+        return row is not None
+
+
+def init_extensions(engine: Engine) -> dict[str, bool]:
+    """按可用性创建扩展，返回各扩展启用状态。"""
+    status: dict[str, bool] = {}
+    for ext in ("btree_gist", "vector", "timescaledb"):
+        available = _extension_available(engine, ext)
+        status[ext] = available
+        if available:
+            with engine.begin() as conn:
+                conn.execute(text(f"CREATE EXTENSION IF NOT EXISTS {ext}"))
+    return status
+
+
+def create_schema(engine: Engine | None = None) -> dict[str, bool]:
+    """创建数据中台全部表，并按需启用时序超表。返回扩展状态。"""
+    engine = engine or get_engine()
+    status = init_extensions(engine)
+    Base.metadata.create_all(engine)
+
+    # 若 TimescaleDB 可用，将时序事件表转为超表（幂等）
+    if status.get("timescaledb"):
+        with engine.begin() as conn:
+            conn.execute(
+                text(
+                    "SELECT create_hypertable('metric_event', 'event_time', "
+                    "if_not_exists => TRUE, migrate_data => TRUE)"
+                )
+            )
+    return status
+
+
+if __name__ == "__main__":
+    st = create_schema()
+    print("数据中台 schema 初始化完成。扩展状态：", st)