freedakgmail
59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
"""数据中台 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)
|