Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from datetime import UTC, datetime | |
| from pathlib import Path | |
| from sqlalchemy import DateTime, Engine, create_engine, event | |
| from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker | |
| from sqlalchemy.types import TypeDecorator | |
| class Base(DeclarativeBase): | |
| pass | |
| class UTCDateTime(TypeDecorator[datetime]): | |
| """Store datetimes as naive UTC and return timezone-aware UTC values.""" | |
| impl = DateTime | |
| cache_ok = True | |
| def process_bind_param(self, value: datetime | None, dialect): | |
| if value is None: | |
| return None | |
| if value.tzinfo is None: | |
| value = value.replace(tzinfo=UTC) | |
| return value.astimezone(UTC).replace(tzinfo=None) | |
| def process_result_value(self, value: datetime | None, dialect): | |
| if value is None: | |
| return None | |
| return value.replace(tzinfo=UTC) | |
| def create_db_engine(database_url: str) -> Engine: | |
| if database_url.startswith("sqlite:///") and ":memory:" not in database_url: | |
| path = database_url.removeprefix("sqlite:///") | |
| Path(path).parent.mkdir(parents=True, exist_ok=True) | |
| connect_args = {"check_same_thread": False} if database_url.startswith("sqlite") else {} | |
| engine = create_engine(database_url, connect_args=connect_args, future=True) | |
| if database_url.startswith("sqlite"): | |
| def _set_sqlite_pragma(dbapi_connection, connection_record): | |
| cursor = dbapi_connection.cursor() | |
| cursor.execute("PRAGMA foreign_keys=ON") | |
| cursor.close() | |
| 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) -> None: | |
| from app import models # noqa: F401 | |
| Base.metadata.create_all(engine) | |