from typing import Callable

from starlette.requests import Request
from starlette.responses import Response

from pillar_tool.db.database import get_connection

async def db_connection_middleware(request: Request, call_next: Callable) -> Response:
    session = get_connection()

    request.state.db = session
    try:
        response: Response = await call_next(request)
        if 200 <= response.status_code <= 299:
            session.commit()
            session.close()
        else:
            session.rollback()
            session.close()
    except Exception as e:
        session.rollback()
        session.close()
        raise e

    return response