File size: 6,823 Bytes
2ac2730 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | """
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
|