Files
sales-trainer/backend/app/db.py

84 lines
2.7 KiB
Python

"""SQLAlchemy engine and declarative metadata foundation for S4 persistence."""
from __future__ import annotations
import os
from pathlib import Path
from typing import Any
from sqlalchemy import create_engine, event
from sqlalchemy.engine import Engine, make_url
from sqlalchemy.orm import DeclarativeBase, sessionmaker
class Base(DeclarativeBase):
"""Base class shared by all relational models."""
SUPPORTED_DATABASE_DIALECTS = frozenset({"sqlite", "postgresql"})
def database_url(configured: str | None = None) -> str:
"""Resolve an explicit URL, then environment, then a local SQLite fallback.
The fallback keeps local schema and migration tests self-contained. Production
deployment should provide ``DATABASE_URL`` pointing at PostgreSQL before the
repository adapters are enabled.
"""
explicit = (configured or "").strip()
if explicit:
return explicit
from .config import Config
env_url = os.environ.get("DATABASE_URL", "").strip() or Config.DATABASE_URL
if env_url:
return env_url
data_dir = Path(os.environ.get("DATA_DIR", str(Config.DATA_DIR))).resolve()
return f"sqlite:///{data_dir / 'sales_trainer.db'}"
def validate_database_url(url: str | None = None) -> str:
"""Resolve a database URL and reject dialects that weaken schema invariants."""
resolved_url = database_url(url)
dialect_name = make_url(resolved_url).get_backend_name()
if dialect_name not in SUPPORTED_DATABASE_DIALECTS:
raise ValueError("Unsupported database dialect; use sqlite or postgresql")
return resolved_url
def create_db_engine(url: str | None = None, **kwargs: Any) -> Engine:
"""Create a SQLite/PostgreSQL engine with SQLite foreign keys enabled."""
resolved_url = validate_database_url(url)
options: dict[str, Any] = {"pool_pre_ping": True, **kwargs}
if resolved_url.startswith("sqlite"):
connect_args = dict(options.pop("connect_args", {}))
connect_args.setdefault("check_same_thread", False)
options["connect_args"] = connect_args
engine = create_engine(resolved_url, **options)
if resolved_url.startswith("sqlite"):
@event.listens_for(engine, "connect")
def _enable_sqlite_foreign_keys(dbapi_connection: Any, _connection_record: Any) -> None:
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
return engine
def create_session_factory(engine: Engine):
"""Return a bounded session factory for a supplied engine."""
return sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
__all__ = [
"Base",
"create_db_engine",
"create_session_factory",
"database_url",
"validate_database_url",
]