Spaces:
Running
Running
| import os | |
| import asyncio | |
| from dotenv import load_dotenv | |
| from sqlalchemy.engine import make_url | |
| from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine | |
| from sqlalchemy.orm import declarative_base | |
| from sqlalchemy.pool import NullPool | |
| # Load environment variables | |
| load_dotenv() | |
| def _resolve_database_url() -> str: | |
| """Return an async-compatible database URL.""" | |
| database_url = os.getenv("DATABASE_URL") | |
| if not database_url: | |
| # Default to SQLite in /data for HF Spaces (mounted bucket) | |
| # Falls back to local SQLite if /data doesn't exist. | |
| db_path = os.getenv("CHROMA_PERSIST_DIR", "/data") | |
| os.makedirs(db_path, exist_ok=True) | |
| database_url = f"sqlite:///{db_path}/rag_lab.db" | |
| print(f"[INFO] DATABASE_URL not set. Using default: {database_url}") | |
| url = make_url(database_url) | |
| drivername = url.drivername.lower() | |
| if drivername.startswith("sqlite") and "+aiosqlite" not in drivername: | |
| return url.set(drivername="sqlite+aiosqlite").render_as_string(hide_password=False) | |
| if drivername.startswith("postgresql") and "+asyncpg" not in drivername: | |
| return url.set(drivername="postgresql+asyncpg").render_as_string(hide_password=False) | |
| return database_url | |
| ASYNC_DATABASE_URL = _resolve_database_url() | |
| # Create async SQLAlchemy engine and session factory. | |
| _engine_kwargs = {"future": True} | |
| # SQLite + async streaming requests can frequently hit cancelled tasks on | |
| # disconnect. Using NullPool avoids reusing cancelled/stale connections. | |
| if ASYNC_DATABASE_URL.startswith("sqlite+aiosqlite"): | |
| _engine_kwargs["poolclass"] = NullPool | |
| # Streaming request cancellations can interrupt rollback/terminate on close. | |
| # Avoid reset-on-return for one-shot SQLite connections to reduce noisy close errors. | |
| _engine_kwargs["pool_reset_on_return"] = None | |
| else: | |
| # Production-friendly pooling for PostgreSQL/other server DBs. | |
| _engine_kwargs.update( | |
| { | |
| "pool_size": int(os.getenv("DB_POOL_SIZE", "10")), | |
| "max_overflow": int(os.getenv("DB_MAX_OVERFLOW", "20")), | |
| "pool_pre_ping": True, | |
| } | |
| ) | |
| engine = create_async_engine(ASYNC_DATABASE_URL, **_engine_kwargs) | |
| AsyncSessionLocal = async_sessionmaker( | |
| bind=engine, | |
| class_=AsyncSession, | |
| autoflush=False, | |
| expire_on_commit=False, | |
| ) | |
| # Create declarative base | |
| Base = declarative_base() | |
| async def get_db(): | |
| """FastAPI dependency for getting an async DB session.""" | |
| async with AsyncSessionLocal() as db: | |
| try: | |
| yield db | |
| finally: | |
| # Client disconnects can cancel request scope while dependency teardown runs. | |
| # Shield close so connection cleanup can complete safely. | |
| await asyncio.shield(db.close()) | |