dekun
72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
"""数据库引擎与会话。SQLite 仅用于本地启动，生产环境把 DATABASE_URL 指到 PostgreSQL。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import DateTime, create_engine
|
|
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
|
|
from sqlalchemy.types import TypeDecorator
|
|
|
|
from app.config import settings
|
|
|
|
|
|
class UtcDateTime(TypeDecorator):
|
|
"""统一读写 UTC，避免 SQLite 把时区信息丢掉后比较失败。"""
|
|
|
|
impl = DateTime
|
|
cache_ok = True
|
|
|
|
def process_bind_param(self, value, dialect):
|
|
if value is None:
|
|
return None
|
|
if value.tzinfo is None:
|
|
value = value.replace(tzinfo=timezone.utc)
|
|
return value.astimezone(timezone.utc)
|
|
|
|
def process_result_value(self, value, dialect):
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, str):
|
|
text = value.replace("Z", "+00:00")
|
|
value = datetime.fromisoformat(text)
|
|
if value.tzinfo is None:
|
|
return value.replace(tzinfo=timezone.utc)
|
|
return value.astimezone(timezone.utc)
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
connect_args = {"check_same_thread": False} if settings.database_url.startswith("sqlite") else {}
|
|
engine = create_engine(settings.database_url, connect_args=connect_args, pool_pre_ping=True)
|
|
SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
|
|
|
|
|
|
@contextmanager
|
|
def session_scope() -> Iterator[Session]:
|
|
session = SessionLocal()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
def get_db() -> Iterator[Session]:
|
|
session = SessionLocal()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|