61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
"""Database engine and session management.
|
|
|
|
Provides a SQLModel engine factory and a session dependency
|
|
for use with FastAPI's dependency injection.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
from typing import Generator
|
|
|
|
import structlog
|
|
from sqlmodel import Session, create_engine
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
|
|
|
|
def get_engine(database_url: str):
|
|
"""Create a SQLAlchemy engine for the given database URL.
|
|
|
|
For SQLite, ensures the parent directory exists and sets
|
|
appropriate connection arguments.
|
|
|
|
Args:
|
|
database_url: SQLAlchemy-compatible connection string.
|
|
|
|
Returns:
|
|
A configured SQLAlchemy Engine.
|
|
"""
|
|
connect_args = {}
|
|
if database_url.startswith("sqlite"):
|
|
# Ensure the data directory exists
|
|
db_path = database_url.replace("sqlite:///", "")
|
|
if db_path and db_path != ":memory:":
|
|
Path(db_path).parent.mkdir(parents=True, exist_ok=True)
|
|
connect_args = {"check_same_thread": False}
|
|
|
|
engine = create_engine(
|
|
database_url,
|
|
connect_args=connect_args,
|
|
echo=False,
|
|
)
|
|
logger.info("database_engine_created", url=database_url)
|
|
return engine
|
|
|
|
|
|
def get_db_session(engine) -> Generator[Session, None, None]:
|
|
"""Yield a SQLModel session for dependency injection.
|
|
|
|
Usage in FastAPI routes:
|
|
@router.get("/example")
|
|
def example(session: Session = Depends(get_db_session)):
|
|
...
|
|
|
|
Args:
|
|
engine: The SQLAlchemy engine to bind the session to.
|
|
|
|
Yields:
|
|
A SQLModel Session that auto-closes after use.
|
|
"""
|
|
with Session(engine) as session:
|
|
yield session
|