A union of curiosity and data science

Knowledgebase and brain dump of a database engineer


Template : FastAPI SQLAlchemy Session Management - Postgres (Dependency injection and transaction management)

from typing import Generator, Optional
from fastapi import Depends, FastAPI, HTTPException
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.exc import SQLAlchemyError
from contextlib import contextmanager

# Database configuration
DATABASE_URL = "postgresql://user:password@localhost/dbname"
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

app = FastAPI()

# Session dependency
def get_db() -> Generator[Session, None, None]:
    db = SessionLocal()
    try:
        yield db
        db.commit()  # Commit if no exception occurred
    except SQLAlchemyError as e:
        db.rollback()  # Rollback on database-related errors
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        db.close()

# Custom transaction context manager for explicit transaction control
@contextmanager
def transaction(db: Session):
    try:
        yield
        db.commit()
    except Exception:
        db.rollback()
        raise

# Example routes demonstrating different transaction patterns

# Pattern 1: Simple single-operation transaction
@app.post("/users/")
def create_user(name: str, db: Session = Depends(get_db)):
    user = User(name=name)
    db.add(user)
    # No need for explicit flush/commit - handled by dependency
    return user

# Pattern 2: Multiple related operations in one transaction
@app.post("/orders/")
def create_order_with_items(
    order_data: OrderCreate,
    db: Session = Depends(get_db)
):
    # Create order and items in a single transaction
    order = Order(customer_id=order_data.customer_id)
    db.add(order)
    db.flush()  # Flush to get the order.id for the items
    
    for item in order_data.items:
        order_item = OrderItem(
            order_id=order.id,
            product_id=item.product_id,
            quantity=item.quantity
        )
        db.add(order_item)
    
    # Commit handled by dependency

# Pattern 3: Multiple independent operations that should be separate transactions
@app.post("/batch-process/")
def batch_process(items: list[ItemData], db: Session = Depends(get_db)):
    results = []
    
    for item in items:
        # Use nested transaction for each item
        with transaction(db):
            processed_item = process_single_item(db, item)
            results.append(processed_item)
    
    return results

# Pattern 4: Complex operation with explicit transaction control
@app.post("/transfer/")
def transfer_funds(
    from_account_id: int,
    to_account_id: int,
    amount: float,
    db: Session = Depends(get_db)
):
    # Explicit transaction with proper error handling
    with transaction(db):
        from_account = db.query(Account).with_for_update().get(from_account_id)
        if not from_account or from_account.balance < amount:
            raise HTTPException(status_code=400, detail="Insufficient funds")
        
        to_account = db.query(Account).with_for_update().get(to_account_id)
        if not to_account:
            raise HTTPException(status_code=404, detail="Recipient account not found")
        
        from_account.balance -= amount
        to_account.balance += amount
        
        db.flush()  # Ensure both updates are valid before commit

# Best practices demonstrated:
# 1. Session management through dependency injection
# 2. Automatic commit/rollback handling
# 3. Proper error handling and propagation
# 4. Strategic use of flush() when needed
# 5. Transaction isolation for batch operations
# 6. Explicit transaction control when needed
# 7. Use of with_for_update() for row-level locking

 

Add comment