database.py

from sqlmodel import SQLModel, create_engine, Session
from typing import Generator
import os

# Import all models to register them with SQLModel metadata
# This ensures all tables are created when init_db() is called
from models import (  # noqa: F401
    Event,
    Notification,
    Wildlife,
    WildlifeStats,
    EventMedia,
    SystemStatus,
)

# SQLite database URL
DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./wildlife.db")

# Create engine
engine = create_engine(DATABASE_URL, echo=False, connect_args={"check_same_thread": False})


def init_db() -> None:
    """Initialize database by creating all tables."""
    SQLModel.metadata.create_all(engine)


def get_session() -> Generator[Session, None, None]:
    """Dependency for getting database session."""
    with Session(engine) as session:
        yield session