28 lines
669 B
Python
28 lines
669 B
Python
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
|
|
|
|
|