TalhaRafique1/Full-stack-Todo

1"""Database connection and session management for SQLModel with Neon PostgreSQL."""2 3import os4from typing import Generator5from sqlmodel import SQLModel, create_engine, Session6 7# Load environment variables from .env file8try:9    from dotenv import load_dotenv10    load_dotenv()11except ImportError:12    pass  # dotenv is optional13 14# Import models to register them with SQLModel metadata15from .models import User, Task  # noqa: F40116 17# Database URL from environment variable18DATABASE_URL = os.getenv(19    "DATABASE_URL",20    "postgresql+psycopg://user:password@host/database"21)22 23# Ensure we're using the psycopg driver24if "postgresql://" in DATABASE_URL:25    DATABASE_URL = DATABASE_URL.replace("postgresql://", "postgresql+psycopg://")26elif "postgres://" in DATABASE_URL:27    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql+psycopg://")28 29# Create SQLModel engine30engine = create_engine(DATABASE_URL, echo=False)31 32 33def get_db() -> Generator[Session, None, None]:34    """Dependency to get database session."""35    with Session(engine) as session:36        yield session37 38 39def init_db() -> None:40    """Initialize database tables."""41    SQLModel.metadata.create_all(engine)42