Spaces:
Sleeping
Sleeping
| """ | |
| Async SQLAlchemy Engine & Session factory for PostgreSQL (Neon) | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import traceback | |
| from typing import Optional | |
| from sqlalchemy.ext.asyncio import ( | |
| AsyncEngine, | |
| AsyncSession, | |
| async_sessionmaker, | |
| create_async_engine, | |
| ) | |
| from config import settings | |
| logger = logging.getLogger(__name__) | |
| engine: Optional[AsyncEngine] = None | |
| async_session_factory: Optional[async_sessionmaker[AsyncSession]] = None | |
| def _build_engine() -> AsyncEngine: | |
| if not settings.database_url: | |
| raise RuntimeError( | |
| "DATABASE_URL is not set. " | |
| "Provide a Neon / PostgreSQL async connection string, e.g. " | |
| "postgresql+asyncpg://user:pass@host/db?sslmode=require" | |
| ) | |
| url = settings.database_url | |
| # Ensure async driver | |
| if url.startswith("postgresql://"): | |
| url = url.replace("postgresql://", "postgresql+asyncpg://", 1) | |
| elif url.startswith("postgres://"): | |
| url = url.replace("postgres://", "postgresql+asyncpg://", 1) | |
| ssl_mode = None | |
| if "?sslmode=" in url: | |
| url_parts = url.split("?") | |
| base_url = url_parts[0] | |
| params = url_parts[1] if len(url_parts) > 1 else "" | |
| param_list = [] | |
| for param in params.split("&"): | |
| if param.startswith("sslmode="): | |
| ssl_mode = param.split("=")[1] | |
| else: | |
| param_list.append(param) | |
| if param_list: | |
| url = f"{base_url}?{'&'.join(param_list)}" | |
| else: | |
| url = base_url | |
| connect_args = {} | |
| if ssl_mode: | |
| if ssl_mode in ("require", "verify-ca", "verify-full", "prefer"): | |
| connect_args["ssl"] = True | |
| elif ssl_mode == "disable": | |
| connect_args["ssl"] = False | |
| elif ssl_mode == "allow": | |
| connect_args["ssl"] = True | |
| if not ssl_mode and "neon.tech" in url: | |
| connect_args["ssl"] = True | |
| logger.info("🔒 Neon detected: enabling SSL by default") | |
| if connect_args: | |
| logger.info(f"🔧 Using connect_args: {connect_args}") | |
| return create_async_engine( | |
| url, | |
| echo=False, | |
| pool_size=10, | |
| max_overflow=20, | |
| pool_pre_ping=True, | |
| pool_recycle=300, | |
| connect_args=connect_args, | |
| ) | |
| async def init_database() -> None: | |
| """Create engine + session factory + tables. Call once at startup.""" | |
| global engine, async_session_factory | |
| if engine is not None: | |
| logger.info("ℹ️ Database already initialized") | |
| return | |
| try: | |
| logger.info("🔧 Building database engine...") | |
| engine = _build_engine() | |
| logger.info("✅ Engine built successfully") | |
| logger.info("🔧 Creating session factory...") | |
| async_session_factory = async_sessionmaker( | |
| bind=engine, | |
| class_=AsyncSession, | |
| expire_on_commit=False, | |
| autoflush=False, | |
| autocommit=False, | |
| ) | |
| logger.info("✅ Session factory created") | |
| from database.base import Base | |
| import database.models # noqa: F401 | |
| logger.info("🔧 Creating tables...") | |
| async with engine.begin() as conn: | |
| await conn.run_sync(Base.metadata.create_all) | |
| logger.info("✅ SQLAlchemy async engine initialized + tables created successfully") | |
| except Exception as e: | |
| logger.error(f"❌ ERROR initializing database: {e}") | |
| logger.error(traceback.format_exc()) | |
| raise | |
| async def close_database() -> None: | |
| """Dispose engine on shutdown.""" | |
| global engine, async_session_factory | |
| if engine is not None: | |
| await engine.dispose() | |
| engine = None | |
| async_session_factory = None | |
| logger.info("✅ SQLAlchemy engine disposed") | |