46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
from collections.abc import Generator
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import create_engine, event
|
|
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
|
|
|
|
from app.config import get_settings
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
def create_sqlite_engine(database_url: str):
|
|
if database_url.startswith("sqlite:///") and database_url != "sqlite:///:memory:":
|
|
db_path = Path(database_url.removeprefix("sqlite:///"))
|
|
if db_path.parent != Path("."):
|
|
db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
connect_args = {"check_same_thread": False, "timeout": 10}
|
|
engine = create_engine(database_url, connect_args=connect_args)
|
|
engine.dialect.connect_args = connect_args
|
|
|
|
@event.listens_for(engine, "connect")
|
|
def set_sqlite_pragma(dbapi_connection, _connection_record):
|
|
cursor = dbapi_connection.cursor()
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.close()
|
|
|
|
return engine
|
|
|
|
|
|
engine = create_sqlite_engine(get_settings().database_url)
|
|
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)
|
|
|
|
|
|
def init_db() -> None:
|
|
import app.models # noqa: F401
|
|
|
|
Base.metadata.create_all(engine)
|
|
|
|
|
|
def get_db_session() -> Generator[Session]:
|
|
with SessionLocal() as session:
|
|
yield session
|