Module 8: Final Project — TaskFlow API
Audit log with triggers + optimistic locking
This capsule adds two features that raise TaskFlow to production-ready: an automatic audit log via PostgreSQL triggers (any change to tasks is recorded in audit.task_log), and optimistic locking with an If-Match header (correct concurrency handling in PUT /tasks/{id}).
By the end, every change to a task is audited with no effort from the app, and two clients editing the same task simultaneously receive correct responses (one succeeds, the other gets 412 with info to resolve).
audit schema and audit.task_log table
alembic revision -m "add audit schema and task_log"# alembic/versions/XXX_add_audit_schema.py
def upgrade() -> None:
# Create the audit schema
op.execute("CREATE SCHEMA IF NOT EXISTS audit")
# audit.task_log table
op.execute("""
CREATE TABLE audit.task_log (
id BIGSERIAL PRIMARY KEY,
task_id UUID NOT NULL,
tenant_id UUID NOT NULL,
user_id UUID, -- who made the change
action TEXT NOT NULL, -- 'INSERT' | 'UPDATE' | 'DELETE'
old_data JSONB, -- previous state (UPDATE/DELETE)
new_data JSONB, -- new state (INSERT/UPDATE)
changed_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)
""")
# Index for common queries
op.execute("""
CREATE INDEX idx_task_log_task_changed
ON audit.task_log(task_id, changed_at DESC)
""")
# Trigger function
op.execute("""
CREATE OR REPLACE FUNCTION audit.log_task_change()
RETURNS TRIGGER AS $$
DECLARE
user_id_setting TEXT;
BEGIN
user_id_setting := current_setting('audit.user_id', TRUE);
IF TG_OP = 'INSERT' THEN
INSERT INTO audit.task_log (
task_id, tenant_id, user_id, action, new_data
)
VALUES (
NEW.id, NEW.tenant_id,
NULLIF(user_id_setting, '')::UUID,
'INSERT',
to_jsonb(NEW)
);
RETURN NEW;
ELSIF TG_OP = 'UPDATE' THEN
INSERT INTO audit.task_log (
task_id, tenant_id, user_id, action, old_data, new_data
)
VALUES (
NEW.id, NEW.tenant_id,
NULLIF(user_id_setting, '')::UUID,
'UPDATE',
to_jsonb(OLD),
to_jsonb(NEW)
);
RETURN NEW;
ELSIF TG_OP = 'DELETE' THEN
INSERT INTO audit.task_log (
task_id, tenant_id, user_id, action, old_data
)
VALUES (
OLD.id, OLD.tenant_id,
NULLIF(user_id_setting, '')::UUID,
'DELETE',
to_jsonb(OLD)
);
RETURN OLD;
END IF;
RETURN NULL;
END;
$$ LANGUAGE plpgsql;
""")
# Trigger on tasks
op.execute("""
CREATE TRIGGER trg_audit_tasks
AFTER INSERT OR UPDATE OR DELETE ON tasks
FOR EACH ROW EXECUTE FUNCTION audit.log_task_change();
""")
def downgrade() -> None:
op.execute("DROP TRIGGER IF EXISTS trg_audit_tasks ON tasks")
op.execute("DROP FUNCTION IF EXISTS audit.log_task_change()")
op.execute("DROP TABLE IF EXISTS audit.task_log")
op.execute("DROP SCHEMA IF EXISTS audit")current_setting('audit.user_id', TRUE) reads the user_id from the session (similar to app.tenant_id). The app must set it before each query.
Pass user_id to the queries
Update the dependency:
# app/deps.py
async def get_db_with_tenant(
tenant_id: str = Depends(get_current_tenant_id),
user_id: str = Depends(get_current_user_id),
) -> AsyncGenerator[AsyncSession, None]:
async with SessionLocal() as session:
await session.execute(
text("SET LOCAL app.tenant_id = :tid"),
{"tid": tenant_id}
)
await session.execute(
text("SET LOCAL audit.user_id = :uid"),
{"uid": user_id}
)
yield session
async def get_current_user_id(
authorization: str = Header(None, alias="Authorization"),
) -> str:
"""Extract user_id from the JWT."""
if not authorization or not authorization.startswith("Bearer "):
raise HTTPException(401, "Missing or invalid Authorization header")
token = authorization.removeprefix("Bearer ")
try:
payload = jwt.decode(token, settings.jwt_secret, algorithms=[settings.jwt_algorithm])
user_id = payload.get("user_id")
if not user_id:
# Fall back to sub if there's no explicit user_id
return payload.get("sub", "")
return user_id
except jwt.PyJWTError:
raise HTTPException(401, "Invalid token")Update auth/login to include user_id in the JWT:
# app/routers/auth.py
@router.post("/auth/login")
async def login(request: LoginRequest):
payload = {
"sub": request.email,
"user_id": str(uuid.uuid4()), # Mock — in production it would come from the DB
"tenant_id": request.tenant_id,
"exp": datetime.now(timezone.utc) + timedelta(hours=24),
}
token = jwt.encode(payload, settings.jwt_secret, algorithm=settings.jwt_algorithm)
return LoginResponse(access_token=token)Endpoint to view a task's history
# app/routers/tasks.py
@router.get("/tasks/{task_id}/history")
async def get_task_history(
task_id: uuid.UUID,
limit: int = Query(50, ge=1, le=200),
db: AsyncSession = Depends(get_db_with_tenant),
):
"""Audit log of a specific task."""
# First verify the task exists (RLS applies, so only for this tenant)
task = await db.scalar(select(Task).where(Task.id == task_id))
if not task:
raise HTTPException(404, "Task not found")
# Query the audit log
result = await db.execute(text("""
SELECT
id, action, old_data, new_data, user_id, changed_at
FROM audit.task_log
WHERE task_id = :task_id
AND tenant_id::text = current_setting('app.tenant_id', TRUE)
ORDER BY changed_at DESC
LIMIT :limit
"""), {"task_id": str(task_id), "limit": limit})
rows = result.mappings().all()
return {"history": [dict(r) for r in rows]}Note the manual WHERE tenant_id::text = current_setting('app.tenant_id', TRUE) on audit.task_log. RLS doesn't apply to the audit schema by default — you'd have to enable it explicitly. For demo simplicity, we do the filter manually.
Optimistic locking on Task
Add a version column:
# app/models/task.py
class Task(Base):
# ... other columns ...
version: Mapped[int] = mapped_column(Integer, default=1, nullable=False)
__mapper_args__ = {
"version_id_col": "version",
}Migration:
alembic revision -m "add version column to tasks"def upgrade() -> None:
op.add_column(
'tasks',
sa.Column('version', sa.Integer, nullable=False, server_default='1')
)
def downgrade() -> None:
op.drop_column('tasks', 'version')PUT /tasks/{id} endpoint with If-Match
# app/routers/tasks.py
from sqlalchemy.orm.exc import StaleDataError
class TaskUpdate(BaseModel):
title: Optional[str] = None
status: Optional[str] = None
@router.put("/tasks/{task_id}", response_model=TaskResponse)
async def update_task(
task_id: uuid.UUID,
data: TaskUpdate,
response: Response,
db: AsyncSession = Depends(get_db_with_tenant),
if_match: Optional[str] = Header(None, alias="If-Match"),
):
if not if_match:
raise HTTPException(
status.HTTP_428_PRECONDITION_REQUIRED,
detail="If-Match header required for updates"
)
try:
requested_version = int(if_match.strip('"'))
except ValueError:
raise HTTPException(400, "Invalid If-Match header format")
# Fetch task (RLS filters by tenant)
task = await db.scalar(
select(Task).where(
Task.id == task_id,
Task.deleted_at.is_(None),
)
)
if not task:
raise HTTPException(404, "Task not found")
# Pre-check
if task.version != requested_version:
raise HTTPException(
status.HTTP_412_PRECONDITION_FAILED,
detail={
"error": "version_mismatch",
"current_etag": f'"{task.version}"',
"your_etag": if_match,
"current_state": {
"id": str(task.id),
"title": task.title,
"status": task.status,
}
}
)
# Apply changes
changes = data.model_dump(exclude_unset=True)
for field, value in changes.items():
setattr(task, field, value)
try:
await db.commit()
except StaleDataError:
await db.rollback()
# Re-fetch for current state
current_task = await db.scalar(
select(Task).where(Task.id == task_id, Task.deleted_at.is_(None))
)
raise HTTPException(
status.HTTP_412_PRECONDITION_FAILED,
detail={
"error": "version_mismatch_race",
"current_etag": f'"{current_task.version}"',
"current_state": {
"id": str(current_task.id),
"title": current_task.title,
"status": current_task.status,
}
}
)
response.headers["ETag"] = f'"{task.version}"'
return TaskResponse(
id=task.id, project_id=task.project_id, title=task.title,
status=task.status, created_at=task.created_at, updated_at=task.updated_at,
)And GET /tasks/{id} must return an ETag:
@router.get("/tasks/{task_id}", response_model=TaskResponse)
async def get_task(
task_id: uuid.UUID,
response: Response,
db: AsyncSession = Depends(get_db_with_tenant),
if_none_match: Optional[str] = Header(None, alias="If-None-Match"),
):
task = await db.scalar(
select(Task).where(Task.id == task_id, Task.deleted_at.is_(None))
)
if not task:
raise HTTPException(404, "Task not found")
current_etag = f'"{task.version}"'
if if_none_match == current_etag:
response.status_code = status.HTTP_304_NOT_MODIFIED
response.headers["ETag"] = current_etag
return None
response.headers["ETag"] = current_etag
return TaskResponse(
id=task.id, project_id=task.project_id, title=task.title,
status=task.status, created_at=task.created_at, updated_at=task.updated_at,
)Tests
# tests/test_audit_log.py
@pytest.mark.asyncio
async def test_create_task_logs_to_audit(client, tenant_a_token, project_a_id, db):
# Create task
response = await client.post(
"/tasks",
json={"project_id": project_a_id, "title": "T", "status": "pending"},
headers={"Authorization": f"Bearer {tenant_a_token}"}
)
task_id = response.json()["id"]
# Verify the audit log
history = await client.get(
f"/tasks/{task_id}/history",
headers={"Authorization": f"Bearer {tenant_a_token}"}
)
log = history.json()["history"]
assert len(log) == 1
assert log[0]["action"] == "INSERT"
assert log[0]["new_data"]["title"] == "T"
@pytest.mark.asyncio
async def test_update_task_logs_to_audit(client, tenant_a_token, project_a_id):
# Create
create = await client.post("/tasks", json={...}, headers={...})
task_id = create.json()["id"]
initial_etag = create.headers.get("ETag")
# Update
await client.put(
f"/tasks/{task_id}",
json={"status": "completed"},
headers={"Authorization": f"Bearer {tenant_a_token}", "If-Match": initial_etag}
)
# History should have 2 entries
history = await client.get(f"/tasks/{task_id}/history", headers={...})
log = history.json()["history"]
assert len(log) == 2
# Ordered DESC
assert log[0]["action"] == "UPDATE"
assert log[1]["action"] == "INSERT"
# Diff captured
assert log[0]["old_data"]["status"] == "pending"
assert log[0]["new_data"]["status"] == "completed"
# tests/test_tasks_optimistic_lock.py
@pytest.mark.asyncio
async def test_concurrent_updates_second_returns_412(client, tenant_a_token, project_a_id):
# Create
create = await client.post("/tasks", json={...}, headers={...})
task_id = create.json()["id"]
etag = create.headers["ETag"]
# Client A updates (success)
response_a = await client.put(
f"/tasks/{task_id}",
json={"title": "By A"},
headers={"Authorization": f"Bearer {tenant_a_token}", "If-Match": etag}
)
assert response_a.status_code == 200
# Client B with the old etag (412)
response_b = await client.put(
f"/tasks/{task_id}",
json={"title": "By B"},
headers={"Authorization": f"Bearer {tenant_a_token}", "If-Match": etag}
)
assert response_b.status_code == 412
assert response_b.json()["detail"]["error"] == "version_mismatch"
@pytest.mark.asyncio
async def test_update_without_if_match_returns_428(client, tenant_a_token, project_a_id):
create = await client.post("/tasks", json={...}, headers={...})
task_id = create.json()["id"]
# No If-Match
response = await client.put(
f"/tasks/{task_id}",
json={"title": "X"},
headers={"Authorization": f"Bearer {tenant_a_token}"}
)
assert response.status_code == 428Pitfalls and common mistakes
1. Trigger without IS DISTINCT FROM in the check.
If you want to avoid logging UPDATEs that don't change anything:
IF TG_OP = 'UPDATE' AND OLD IS DISTINCT FROM NEW THEN
-- log
END IF;Without this, every PUT (even if the values are the same) creates an audit entry.
2. current_setting('audit.user_id') without NULL handling.
If the app doesn't set the setting, NULLIF(...)::UUID can fail. The NULLIF(value, '') converts an empty string to NULL before the cast.
3. The audit schema without RLS.
If you want to filter by tenant in the audit log, RLS there too. For demo simplicity, we do the filter manually in the query.
4. A trigger that fails silently.
If the audit.log_task_change() trigger fails, the INSERT/UPDATE/DELETE on tasks also fails (because it's an AFTER trigger inside the transaction). Important to handle the trigger's errors gracefully.
5. to_jsonb(NEW) with large columns.
to_jsonb serializes the whole row. If you have large columns (TEXT, BLOB), the audit grows fast. Consider to_jsonb(NEW) - 'large_column' to exclude them.
6. Pre-check + post-check redundancy.
The pre-check (in Python code) and the post-check (SQLAlchemy's StaleDataError) cover different cases:
- Pre-check: fast, before modifying.
- Post-check: catches a race condition between the pre-check and the commit.
Both are necessary.
7. Response without an ETag header.
Without response.headers["ETag"], the client can't maintain the version state. ETag mandatory in GET and PUT responses.
Summary and next step
What you have now:
- The
auditschema with atask_logtable populated by triggers. - A PostgreSQL trigger that captures INSERT/UPDATE/DELETE on
tasks. - Context variables:
app.tenant_idandaudit.user_id. - A history endpoint:
GET /tasks/{id}/history. - A
versioncolumn with__mapper_args__. PUT /tasks/{id}withIf-Match→ 412 /StaleDataError→ 412.- Tests for the audit log and optimistic locking.
Commit:
git add .
git commit -m "feat: audit log via triggers + optimistic locking with If-Match"In the next capsule we close out the patterns with two peak features: a bulk endpoint with COPY + ON CONFLICT and a zero-downtime migration executed live with wrk running. That's the module's binary criterion: 0 errors in 600s of traffic during the 3-phase deploy.
Resources
- PostgreSQL — Trigger Functions — reference.
- PostgreSQL —
to_jsonb— serialization. - SQLAlchemy 2.0 — Versioning Counter — reference.
- RFC 7232 — Conditional Requests —
If-Match. - Audit triggers in PostgreSQL — patterns reference.
Capsule 05 of 08 — Module 8 — SQL Patterns for Production APIs Guide