72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
import uuid
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import Depends
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from src.config import settings
|
|
from src.db import SessionLocal, request_context
|
|
|
|
|
|
async def get_db():
|
|
"""FastAPI Dependency Generator for Database"""
|
|
|
|
context = request_context.get() or "unknown"
|
|
|
|
db: AsyncSession = SessionLocal()
|
|
try:
|
|
if settings.DB.TRACING:
|
|
await db.execute(
|
|
text("SELECT set_config('application_name', :name, false)"),
|
|
{"name": context},
|
|
)
|
|
yield db
|
|
except Exception:
|
|
await db.rollback()
|
|
raise
|
|
finally:
|
|
# Always send ROLLBACK unconditionally so the wire-level transaction
|
|
# is closed before the TCP connection drops. Supavisor v2 does NOT
|
|
# clean up orphaned transactions on client disconnect in transaction-
|
|
# pooling mode, so relying on `in_transaction()` (Python-side state)
|
|
# can leave the backend pinned with an open BEGIN.
|
|
await db.rollback()
|
|
await db.close()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def tracked_db(operation_name: str | None = None):
|
|
"""Context manager for tracked database sessions"""
|
|
# Get request ID if available, or create operation-specific one
|
|
context = request_context.get()
|
|
token = None
|
|
|
|
if not context and operation_name:
|
|
context = f"task:{operation_name}:{str(uuid.uuid4())[:8]}"
|
|
token = request_context.set(context)
|
|
|
|
# Create session with tracking info
|
|
db = SessionLocal()
|
|
|
|
try:
|
|
if settings.DB.TRACING:
|
|
await db.execute(
|
|
text("SELECT set_config('application_name', :name, false)"),
|
|
{"name": context or f"task:{operation_name}"},
|
|
)
|
|
|
|
yield db
|
|
except Exception:
|
|
await db.rollback()
|
|
raise
|
|
finally:
|
|
# Always send ROLLBACK unconditionally — see get_db() comment.
|
|
await db.rollback()
|
|
await db.close()
|
|
if token: # Only reset if we set it
|
|
request_context.reset(token)
|
|
|
|
|
|
db: AsyncSession = Depends(get_db)
|