gisaf-backend/src/gisaf/database.py

from contextlib import asynccontextmanager
from typing import Annotated
from collections.abc import AsyncGenerator

from sqlalchemy.ext.asyncio import create_async_engine
from sqlmodel.ext.asyncio.session import AsyncSession
from fastapi import Depends
import pandas as pd

from .config import conf

engine = create_async_engine(
    conf.db.get_sqla_url(),
    echo=conf.db.echo,
    pool_size=conf.db.pool_size,
    max_overflow=conf.db.max_overflow,
)

async def get_db_session() -> AsyncGenerator[AsyncSession]:
    async with AsyncSession(engine) as session:
        yield session

@asynccontextmanager
async def db_session() -> AsyncGenerator[AsyncSession]:
    async with AsyncSession(engine) as session:
        yield session

def pandas_query(session, query):
    return pd.read_sql_query(query, session.connection())

fastapi_db_session = Annotated[AsyncSession, Depends(get_db_session)]