""" Database connection management for Todo AI Chatbot. Provides SQLModel engine creation, connection verification, and table management for Neon PostgreSQL. """ import logging from typing import Generator from sqlmodel import create_engine, text, SQLModel, Session from sqlalchemy import Engine from .config import get_settings logger = logging.getLogger(__name__) # Global engine instance _engine: Engine | None = None def get_engine() -> Engine: """ Get or create the SQLModel database engine. Uses postgresql+psycopg driver for Neon serverless PostgreSQL. SSL is enabled automatically via sslmode=require in connection string. Returns: SQLAlchemy Engine configured for Neon PostgreSQL. """ global _engine if _engine is None: settings = get_settings() # Convert postgres:// to postgresql+psycopg:// for SQLAlchemy 2.0 database_url = settings.database_url if database_url.startswith("postgres://"): database_url = database_url.replace("postgres://", "postgresql+psycopg://", 1) elif database_url.startswith("postgresql://"): database_url = database_url.replace("postgresql://", "postgresql+psycopg://", 1) _engine = create_engine( database_url, echo=False, # Set to True for SQL query logging pool_pre_ping=True, # Verify connections before use ) logger.info("Database engine created successfully") return _engine def check_connection() -> bool: """ Check if the database connection is working. Executes a simple query to verify connectivity. Returns: True if connection is successful, False otherwise. """ try: engine = get_engine() with engine.connect() as conn: conn.execute(text("SELECT 1")) conn.commit() logger.info("Database connection verified") return True except Exception as e: logger.error(f"Database connection failed: {e}") return False def close_engine() -> None: """Close the database engine and release connections.""" global _engine if _engine is not None: _engine.dispose() _engine = None logger.info("Database engine closed") def create_tables() -> None: """ Create all database tables from registered SQLModel models. This function uses SQLModel.metadata.create_all() to create tables for all models that have been imported and registered. Note: Models must be imported before calling this function for their tables to be created. """ # Import models to register them with SQLModel.metadata # This ensures the models are known before create_all is called from . import models # noqa: F401 engine = get_engine() try: SQLModel.metadata.create_all(engine) logger.info("Database tables created successfully") except Exception as e: logger.error(f"Failed to create database tables: {e}") raise def get_session() -> Generator[Session, None, None]: """ Get a database session for dependency injection in FastAPI. Yields: SQLModel Session for database operations. """ with Session(get_engine()) as session: yield session def migrate_tasks_table() -> None: """ Add missing columns to tasks table if they don't exist. This handles schema migration for priority and category columns. """ engine = get_engine() try: with engine.connect() as conn: # Check if priority column exists result = conn.execute(text(""" SELECT column_name FROM information_schema.columns WHERE table_name = 'tasks' AND column_name = 'priority' """)) if not result.fetchone(): conn.execute(text("ALTER TABLE tasks ADD COLUMN priority VARCHAR(10)")) logger.info("Added 'priority' column to tasks table") # Check if category column exists result = conn.execute(text(""" SELECT column_name FROM information_schema.columns WHERE table_name = 'tasks' AND column_name = 'category' """)) if not result.fetchone(): conn.execute(text("ALTER TABLE tasks ADD COLUMN category VARCHAR(50)")) logger.info("Added 'category' column to tasks table") # Check if due_date column exists result = conn.execute(text(""" SELECT column_name FROM information_schema.columns WHERE table_name = 'tasks' AND column_name = 'due_date' """)) if not result.fetchone(): conn.execute(text("ALTER TABLE tasks ADD COLUMN due_date TIMESTAMP")) logger.info("Added 'due_date' column to tasks table") conn.commit() logger.info("Tasks table migration completed") except Exception as e: logger.error(f"Failed to migrate tasks table: {e}") raise def migrate_users_table() -> None: """ Add missing columns to users table if they don't exist. """ engine = get_engine() try: with engine.connect() as conn: # Check if password_hash column exists result = conn.execute(text(""" SELECT column_name FROM information_schema.columns WHERE table_name = 'users' AND column_name = 'password_hash' """)) if not result.fetchone(): conn.execute(text("ALTER TABLE users ADD COLUMN password_hash VARCHAR(255)")) logger.info("Added 'password_hash' column to users table") # Check if provider column exists result = conn.execute(text(""" SELECT column_name FROM information_schema.columns WHERE table_name = 'users' AND column_name = 'provider' """)) if not result.fetchone(): conn.execute(text("ALTER TABLE users ADD COLUMN provider VARCHAR(50) DEFAULT 'credentials'")) logger.info("Added 'provider' column to users table") # Check if provider_id column exists result = conn.execute(text(""" SELECT column_name FROM information_schema.columns WHERE table_name = 'users' AND column_name = 'provider_id' """)) if not result.fetchone(): conn.execute(text("ALTER TABLE users ADD COLUMN provider_id VARCHAR(255)")) logger.info("Added 'provider_id' column to users table") conn.commit() logger.info("Users table migration completed") except Exception as e: logger.error(f"Failed to migrate users table: {e}") raise