Spaces:
Sleeping
Sleeping
File size: 2,642 Bytes
fee2f22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """Database engine and session management."""
from __future__ import annotations
import uuid
from typing import Generator
from sqlalchemy import JSON, Column, DateTime, String, Table, create_engine, func
from sqlalchemy.dialects.postgresql import UUID as PG_UUID
from sqlalchemy.orm import declarative_base, sessionmaker
from sqlalchemy.types import CHAR, TypeDecorator
from backend.config import settings
class GUID(TypeDecorator):
"""Platform-independent GUID type that uses CHAR(32) on SQLite and UUID on PostgreSQL."""
impl = CHAR
cache_ok = True
def load_dialect_impl(self, dialect):
"""Adjust based on dialect."""
if dialect.name == "postgresql":
return dialect.type_descriptor(PG_UUID(as_uuid=True))
else:
return dialect.type_descriptor(CHAR(32))
def process_bind_param(self, value, dialect):
"""Convert UUID to string for storage."""
if value is None:
return value
elif dialect.name == "postgresql":
return str(value)
else:
if isinstance(value, uuid.UUID):
return value.hex
else:
return value
def process_result_value(self, value, dialect):
"""Convert string back to UUID."""
if value is None:
return value
else:
if isinstance(value, uuid.UUID):
return value
else:
return uuid.UUID(value)
engine = create_engine(
settings.database_url,
pool_pre_ping=True,
echo=False,
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
model_versions = Table(
"model_versions",
Base.metadata,
Column("id", GUID(), primary_key=True, default=uuid.uuid4),
Column("model_name", String(255), nullable=False, index=True),
Column("version", String(50), nullable=False, index=True),
Column("metrics", JSON, nullable=True),
Column("created_at", DateTime(timezone=True), server_default=func.now(), nullable=False),
)
def init_db() -> None:
"""Create all metadata-defined tables."""
# Import ORM models so SQLAlchemy registers their metadata before create_all.
from backend.models.audit_log import AuditLog # noqa: F401
from backend.models.cleaning_job import CleaningJob # noqa: F401
from backend.models.dataset import Dataset # noqa: F401
Base.metadata.create_all(bind=engine)
def get_db() -> Generator:
"""Yield database session for request-scoped usage."""
db = SessionLocal()
try:
yield db
finally:
db.close()
|