File size: 1,654 Bytes
bef5e76 | 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 | """Async PostgreSQL connection management."""
from sqlalchemy.engine import make_url
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import declarative_base
from src.config.settings import settings
# asyncpg doesn't support libpq query params like sslmode/channel_binding.
# Use SQLAlchemy's URL parser to strip all query params cleanly, then pass ssl via connect_args.
_url = make_url(settings.postgres_connstring).set(drivername="postgresql+asyncpg", query={})
# Separate asyncpg engine for PGVector with prepared_statement_cache_size=0.
# PGVector runs advisory_lock + CREATE EXTENSION as a single multi-statement string.
# asyncpg normally uses prepared statements which reject multi-statement SQL.
# Setting cache_size=0 forces asyncpg to use execute() instead of prepare(),
# which supports multiple statements — no psycopg3 needed, no ProactorEventLoop issue.
_pgvector_engine = create_async_engine(
_url,
pool_pre_ping=True,
connect_args={
"ssl": "require",
"prepared_statement_cache_size": 0,
},
)
engine = create_async_engine(
_url,
echo=False,
pool_pre_ping=True,
pool_size=5,
max_overflow=10,
connect_args={"ssl": "require"},
)
AsyncSessionLocal = async_sessionmaker(
engine,
class_=AsyncSession,
expire_on_commit=False,
autocommit=False,
autoflush=False
)
Base = declarative_base()
async def get_db():
"""Get database session dependency for FastAPI."""
async with AsyncSessionLocal() as session:
try:
yield session
finally:
await session.close()
|