from __future__ import annotations
import os
from sqlalchemy import Engine, create_engine, event, inspect, text
from sqlalchemy.orm import Session, sessionmaker
from witty_service.persistence.orm import Base
def create_sqlite_engine(database_url: str) -> Engine:
engine = create_engine(
database_url,
connect_args={"check_same_thread": False},
future=True,
)
_enable_sqlite_foreign_keys(engine)
return engine
def create_session_factory(engine: Engine) -> sessionmaker[Session]:
return sessionmaker(bind=engine, expire_on_commit=False, class_=Session)
def init_db(engine: Engine, *, auto_create: bool | None = None) -> None:
"""Bootstrap schema only for tests or explicitly opted-in local setups."""
if auto_create is None:
auto_create = os.getenv("WITTY_DATABASE_AUTO_CREATE", "false").lower() in (
"1",
"true",
"yes",
)
if not auto_create:
return
_migrate_models_table(engine)
Base.metadata.create_all(bind=engine)
def _migrate_models_table(engine: Engine) -> None:
"""自动迁移 models 表,添加缺失的 compatibility 列"""
inspector = inspect(engine)
if inspector.has_table("models"):
columns = inspector.get_columns("models")
column_names = {col["name"] for col in columns}
if "compatibility" not in column_names:
with engine.connect() as conn:
conn.execute(text("ALTER TABLE models ADD COLUMN compatibility VARCHAR(16)"))
conn.commit()
def _enable_sqlite_foreign_keys(engine: Engine) -> None:
@event.listens_for(engine, "connect")
def _set_sqlite_pragma(dbapi_connection, _connection_record) -> None:
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()