from contextlib import asynccontextmanager
from typing import Annotated

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

echo = False
pg_url = "postgresql+asyncpg://avgis@localhost/avgis"

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

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

@asynccontextmanager
async def db_session() -> 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)]