51 lines
1.3 KiB
Python
51 lines
1.3 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
from app.core.config import settings
|
|
|
|
_connect_args: dict[str, object] = {}
|
|
_engine_kwargs: dict[str, object] = {
|
|
"pool_pre_ping": True,
|
|
"future": True,
|
|
"pool_size": 5,
|
|
"max_overflow": 10,
|
|
"pool_recycle": 1800,
|
|
}
|
|
|
|
if settings.database_url.startswith("sqlite"):
|
|
_connect_args["check_same_thread"] = False
|
|
_engine_kwargs["poolclass"] = StaticPool
|
|
_engine_kwargs.pop("pool_size", None)
|
|
_engine_kwargs.pop("max_overflow", None)
|
|
_engine_kwargs.pop("pool_recycle", None)
|
|
|
|
engine = create_engine(settings.database_url, connect_args=_connect_args, **_engine_kwargs)
|
|
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True)
|
|
|
|
|
|
@contextmanager
|
|
def session_scope() -> Generator[Session, None, None]:
|
|
session = SessionLocal()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
def get_db() -> Generator[Session, None, None]:
|
|
session = SessionLocal()
|
|
try:
|
|
yield session
|
|
finally:
|
|
session.close()
|