feat: add database engine and session management
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
60
app/database.py
Normal file
60
app/database.py
Normal file
@@ -0,0 +1,60 @@
|
||||
"""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
|
||||
26
tests/test_database.py
Normal file
26
tests/test_database.py
Normal file
@@ -0,0 +1,26 @@
|
||||
"""Tests for database engine and session management."""
|
||||
|
||||
from sqlmodel import Session
|
||||
|
||||
from app.database import get_engine, get_db_session
|
||||
|
||||
|
||||
class TestDatabase:
|
||||
"""Tests for database connection management."""
|
||||
|
||||
def test_get_engine_returns_engine(self) -> None:
|
||||
"""get_engine should return a SQLAlchemy engine instance."""
|
||||
engine = get_engine("sqlite:///data/test.db")
|
||||
assert engine is not None
|
||||
assert "sqlite" in str(engine.url)
|
||||
|
||||
def test_get_db_session_yields_session(self) -> None:
|
||||
"""get_db_session should yield a usable SQLModel Session."""
|
||||
engine = get_engine("sqlite:///:memory:")
|
||||
session_gen = get_db_session(engine)
|
||||
session = next(session_gen)
|
||||
assert isinstance(session, Session)
|
||||
try:
|
||||
next(session_gen)
|
||||
except StopIteration:
|
||||
pass
|
||||
Reference in New Issue
Block a user