1  """
 2  Database connection setup using SQLAlchemy with async support.
 3  """
 4  from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
 5  from sqlalchemy.orm import DeclarativeBase
 6  from config import settings
 7  
 8  
 9  # Create async engine
10  # Note: asyncpg uses postgresql+asyncpg:// instead of postgresql://
11  database_url = settings.database_url.replace("postgresql://", "postgresql+asyncpg://")
12  
13  engine = create_async_engine(
14      database_url,
15      echo=True,  # Set to False in production - logs all SQL queries
16      future=True
17  )
18  
19  # Session factory
20  AsyncSessionLocal = async_sessionmaker(
21      engine,
22      class_=AsyncSession,
23      expire_on_commit=False
24  )
25  
26  
27  # Base class for all models
28  class Base(DeclarativeBase):
29      pass
30  
31  
32  # Dependency to get database session
33  async def get_db() -> AsyncSession:
34      """
35      Dependency that provides a database session.
36      
37      Usage in route:
38          @router.get("/example")
39          async def example(db: AsyncSession = Depends(get_db)):
40              # use db here
41      """
42      async with AsyncSessionLocal() as session:
43          try:
44              yield session
45              await session.commit()
46          except Exception:
47              await session.rollback()
48              raise
49          finally:
50              await session.close()
51