Spaces:
Runtime error
Runtime error
| import os | |
| from typing import AsyncGenerator | |
| import ssl | |
| from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker | |
| from sqlalchemy.orm import declarative_base | |
| # Get DATABASE_URL and ensure it uses asyncpg | |
| DATABASE_URL = os.getenv("DATABASE_URL", "") | |
| # Force asyncpg driver if not specified | |
| if DATABASE_URL.startswith("postgresql://"): | |
| DATABASE_URL = DATABASE_URL.replace("postgresql://", "postgresql+asyncpg://", 1) | |
| # Remove sslmode parameter if present (asyncpg doesn't support it) | |
| if "?sslmode=" in DATABASE_URL: | |
| DATABASE_URL = DATABASE_URL.split("?sslmode=")[0] | |
| # Create SSL context for asyncpg | |
| ssl_context = ssl.create_default_context() | |
| ssl_context.check_hostname = False | |
| ssl_context.verify_mode = ssl.CERT_NONE | |
| engine = create_async_engine( | |
| DATABASE_URL, | |
| echo=True, | |
| future=True, | |
| connect_args={ | |
| "ssl": ssl_context, | |
| }, | |
| ) | |
| AsyncSessionLocal = async_sessionmaker( | |
| engine, | |
| class_=AsyncSession, | |
| expire_on_commit=False, | |
| ) | |
| Base = declarative_base() | |
| async def get_session() -> AsyncGenerator[AsyncSession, None]: | |
| """Get database session.""" | |
| async with AsyncSessionLocal() as session: | |
| yield session | |
| async def init_db() -> None: | |
| """Initialize database - create all tables.""" | |
| async with engine.begin() as conn: | |
| await conn.run_sync(Base.metadata.create_all) | |