- Add database.py with SQLAlchemy engine and session - Implement get_db() for FastAPI dependency injection - Implement init_db() for table creation - Use SQLAlchemy 2.0 declarative_base() syntax - Add comprehensive tests with 100% coverage Tests: 11 passed, 100% coverage
68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
"""Database connection and session management.
|
|
|
|
T06: Database connection & session management
|
|
"""
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker, Session, declarative_base
|
|
from typing import Generator
|
|
|
|
from openrouter_monitor.config import get_settings
|
|
|
|
|
|
# Create declarative base for models (SQLAlchemy 2.0 style)
|
|
Base = declarative_base()
|
|
|
|
# Get settings
|
|
settings = get_settings()
|
|
|
|
# Create engine with SQLite configuration
|
|
# check_same_thread=False is required for SQLite with async/threads
|
|
engine = create_engine(
|
|
settings.database_url,
|
|
connect_args={"check_same_thread": False}
|
|
)
|
|
|
|
# Create session maker with expire_on_commit=False
|
|
# This prevents attributes from being expired after commit
|
|
SessionLocal = sessionmaker(
|
|
autocommit=False,
|
|
autoflush=False,
|
|
bind=engine,
|
|
expire_on_commit=False
|
|
)
|
|
|
|
|
|
def get_db() -> Generator[Session, None, None]:
|
|
"""Get database session for FastAPI dependency injection.
|
|
|
|
This function creates a new database session and yields it.
|
|
The session is automatically closed when the request is done.
|
|
|
|
Yields:
|
|
Session: SQLAlchemy database session
|
|
|
|
Example:
|
|
>>> from fastapi import Depends
|
|
>>> @app.get("/items/")
|
|
>>> def read_items(db: Session = Depends(get_db)):
|
|
... return db.query(Item).all()
|
|
"""
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def init_db() -> None:
|
|
"""Initialize database by creating all tables.
|
|
|
|
This function creates all tables registered with Base.metadata.
|
|
Should be called at application startup.
|
|
|
|
Example:
|
|
>>> from openrouter_monitor.database import init_db
|
|
>>> init_db() # Creates all tables
|
|
"""
|
|
Base.metadata.create_all(bind=engine)
|