Spaces:
Sleeping
Sleeping
File size: 3,830 Bytes
251964d | 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 | """
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")
|