Spaces:
Running
Running
| 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() | |