from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker, declarative_base from config import settings DATABASE_URL = settings.database_url if not DATABASE_URL: raise ValueError( "DATABASE_URL not found in environment variables. " "Please set it in .env file. Example: " "postgresql+asyncpg://user:password@localhost:5432/kairo_db" ) # Create async engine engine = create_async_engine( DATABASE_URL, echo=False, # Set to True for SQL logging future=True, pool_pre_ping=True, # Verify connections before using pool_size=10, max_overflow=20, ) # Create async session factory async_session = sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, autoflush=False ) # Base class for all models Base = declarative_base() async def get_db() -> AsyncSession: """Dependency for getting async database session""" async with async_session() as session: try: yield session finally: await session.close() async def init_db(): """Initialize database by creating all tables""" async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) # Safely migrate existing Postgres table to add the 'what' column try: from sqlalchemy import text await conn.execute(text("ALTER TABLE portfolio_items ADD COLUMN IF NOT EXISTS what TEXT;")) print("[Migration] 'what' column migration completed successfully") except Exception as e: print(f"[Migration Warning] SQL migration warning (might be SQLite or column already exists): {e}") # Safely migrate existing Postgres table to add the 'sort_order' column try: from sqlalchemy import text await conn.execute(text("ALTER TABLE portfolio_items ADD COLUMN IF NOT EXISTS sort_order INTEGER DEFAULT 0 NOT NULL;")) print("[Migration] 'sort_order' column migration completed successfully") except Exception as e: print(f"[Migration Warning] SQL migration warning for sort_order: {e}") # Safely migrate existing Postgres table to add visitor tracking columns try: from sqlalchemy import text await conn.execute(text("ALTER TABLE visitor_telemetry ADD COLUMN IF NOT EXISTS visitor_id VARCHAR(100);")) await conn.execute(text("ALTER TABLE visitor_telemetry ADD COLUMN IF NOT EXISTS session_id VARCHAR(100);")) await conn.execute(text("ALTER TABLE visitor_telemetry ADD COLUMN IF NOT EXISTS ip_address VARCHAR(50);")) await conn.execute(text("ALTER TABLE visitor_telemetry ADD COLUMN IF NOT EXISTS user_agent TEXT;")) await conn.execute(text("ALTER TABLE visitor_telemetry ADD COLUMN IF NOT EXISTS country VARCHAR(100);")) await conn.execute(text("ALTER TABLE visitor_telemetry ADD COLUMN IF NOT EXISTS city VARCHAR(100);")) await conn.execute(text("ALTER TABLE visitor_telemetry ADD COLUMN IF NOT EXISTS visit_count INTEGER DEFAULT 1;")) print("[Migration] VisitorTelemetry schema migrations completed successfully!") except Exception as e: print(f"[Migration Warning] SQL migration warning for visitor_telemetry: {e}") async def close_db(): """Close database connections""" await engine.dispose()