""" Database Streaming Example - Thread-safe SQLAlchemy session management in SSE This example demonstrates: - CORRECT: Create database sessions within generators (thread-safe) - INCORRECT: Reusing sessions across task boundaries (thread-unsafe) - Streaming query results with proper async session management - Clean separation of concerns between endpoint logic and data streaming CRITICAL SESSION RULE: Never reuse database sessions between async tasks. Always create new sessions within generators that will run in task groups. Usage: python 03_database_streaming.py Test with curl: # Stream all tasks curl -N http://localhost:8000/tasks/stream # Stream with filters curl -N http://localhost:8000/tasks/stream?completed=true curl -N http://localhost:8000/tasks/stream?completed=false curl -N http://localhost:8000/tasks/stream?limit=3 # Compare with regular JSON endpoint curl http://localhost:8000/tasks/list """ import asyncio import typing as T from typing import AsyncGenerator, Optional import sqlalchemy as sa import uvicorn from fastapi import Depends, FastAPI, Query from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from starlette.requests import Request from sse_starlette import EventSourceResponse db_bind = create_async_engine("sqlite+aiosqlite:///:memory:") AsyncSessionLocal = async_sessionmaker(bind=db_bind, expire_on_commit=False) # Dependency for regular endpoints (NOT for use in generators) async def async_db_session(): """ Standard dependency for regular endpoints. IMPORTANT: Do NOT use this dependency in SSE generators. The session created here runs in the main request context, but generators run in separate tasks within anyio task groups. """ async with AsyncSessionLocal() as session: yield session AsyncDbSessionDependency = T.Annotated[AsyncSession, Depends(async_db_session)] TODOS_CTE_SQL = """ WITH todo AS (SELECT 1 AS id, 'Implement SSE streaming' AS title, 'Add server-sent events to API' AS description, 0 AS completed UNION ALL SELECT 2, 'Database integration', 'Connect SQLAlchemy with SSE', 1 UNION ALL SELECT 3, 'Add filtering', 'Support query parameters in streams', 0 UNION ALL SELECT 4, 'Write documentation', 'Document best practices', 1 UNION ALL SELECT 5, 'Performance testing', 'Test with large datasets', 0) SELECT * FROM todo """ app = FastAPI() class TaskDatabaseStream: """ Stream that yields database query results as SSE events. Key design decisions: 1. Create dedicated session within the generator (thread-safe) 2. Apply filters inside the database query (efficient) 3. Implement proper async iteration protocol 4. Handle client disconnection gracefully """ def __init__(self, request: Request, completed_filter: Optional[bool] = None, limit: Optional[int] = None): self.request = request self.completed_filter = completed_filter self.limit = limit def __aiter__(self) -> "TaskDatabaseStream": return self async def __anext__(self) -> dict: """ Lazy evaluation: Database query happens here, not in __init__. This ensures the query runs within the correct async context and uses a fresh session created specifically for this stream. """ if not hasattr(self, '_results'): await self._execute_query() self._index = 0 if self._index >= len(self._results): raise StopAsyncIteration if await self.request.is_disconnected(): raise StopAsyncIteration result = self._results[self._index] self._index += 1 # Add small delay to demonstrate streaming behavior await asyncio.sleep(0.3) return { "data": { "id": result.id, "title": result.title, "description": result.description, "completed": bool(result.completed) }, "event": "task", "id": str(result.id) } async def _execute_query(self): """ CORRECT PATTERN: Create session within the generator context. This session is created in the same async context where the generator will yield results, ensuring thread safety with anyio task groups used by EventSourceResponse. Always create new sessions within generators, never reuse sessions from dependency injection. """ async with AsyncSessionLocal() as session: # Build query with filters - using sa.text() for raw SQL query = TODOS_CTE_SQL if self.completed_filter is not None: query += " WHERE completed = :completed" query += " ORDER BY id" if self.limit: query += " LIMIT :limit" # Prepare parameters for SQLAlchemy params = {} if self.completed_filter is not None: params['completed'] = int(self.completed_filter) if self.limit: params['limit'] = self.limit # Execute query using async iteration pattern from example result = await session.execute(sa.text(query), params) self._results = result.fetchall() async def correct_task_stream(request: Request, completed: Optional[bool] = None, limit: Optional[int] = None) -> AsyncGenerator[dict, None]: """ CORRECT: Alternative implementation using async generator function. This demonstrates the same thread-safe pattern using a function instead of a class. The key principle remains: create the session within the generator scope. "Do *NOT* reuse db_session here within the AsyncGenerator, create a new session instead." """ # Session created INSIDE the generator - this is safe async with AsyncSessionLocal() as session: query = TODOS_CTE_SQL params = {} if completed is not None: query += " WHERE completed = :completed" params['completed'] = int(completed) query += " ORDER BY id" if limit: query += " LIMIT :limit" params['limit'] = limit # Execute query and iterate over results result = await session.execute(sa.text(query), params) rows = result.fetchall() for row in rows: if await request.is_disconnected(): break yield { "data": { "id": row.id, "title": row.title, "description": row.description, "completed": bool(row.completed) }, "event": "task", "id": str(row.id) } await asyncio.sleep(0.3) async def incorrect_task_stream_example(session: AsyncSession, request: Request) -> AsyncGenerator[dict, None]: """ INCORRECT: This pattern will cause issues with anyio task groups. The session parameter comes from the dependency injection system, which runs in the main request context. However, this generator will be consumed by EventSourceResponse within an anyio task group, creating a cross-task session usage that can cause threading issues. DO NOT USE THIS PATTERN - included only for educational purposes. """ # This query would fail because session is from different async context result = await session.execute(sa.text(TODOS_CTE_SQL)) rows = result.fetchall() for row in rows: yield {"data": dict(row._mapping), "event": "task"} @app.get("/tasks/stream") async def stream_tasks_endpoint( request: Request, completed: Optional[bool] = Query(None, description="Filter by completion status"), limit: Optional[int] = Query(None, description="Limit number of results") ) -> EventSourceResponse: """ SSE endpoint that streams database query results. Notice: We do NOT inject a database session here because the session must be created within the generator context for thread safety. """ stream = TaskDatabaseStream(request, completed, limit) return EventSourceResponse(stream) @app.get("/tasks/stream-function") async def stream_tasks_function_endpoint( request: Request, completed: Optional[bool] = Query(None, description="Filter by completion status"), limit: Optional[int] = Query(None, description="Limit number of results") ) -> EventSourceResponse: """ Alternative SSE endpoint using generator function instead of class. Both approaches are valid - choose based on complexity of your logic. Classes are better for stateful streaming, functions for simple cases. """ return EventSourceResponse(correct_task_stream(request, completed, limit)) @app.get("/tasks/list") async def list_tasks_endpoint( session: AsyncDbSessionDependency, completed: Optional[bool] = Query(None, description="Filter by completion status"), limit: Optional[int] = Query(None, description="Limit number of results") ) -> dict: """ Regular JSON endpoint - safe to use dependency injection. This demonstrates the CORRECT use of session dependency for regular endpoints vs the INCORRECT use in SSE generators. """ query = TODOS_CTE_SQL params = {} if completed is not None: query += " WHERE completed = :completed" params['completed'] = int(completed) query += " ORDER BY id" if limit: query += " LIMIT :limit" params['limit'] = limit result = await session.execute(sa.text(query), params) rows = result.fetchall() return { "tasks": [dict(row._mapping) for row in rows], "total": len(rows), "streaming_available": True } @app.get("/status") async def status_endpoint(): """Simple status endpoint.""" return {"status": "ready", "database": "in-memory", "streaming": "enabled"} if __name__ == "__main__": print("Database Streaming SSE Server") print("Stream all: curl -N http://localhost:8000/tasks/stream") print("Stream filtered: curl -N http://localhost:8000/tasks/stream?completed=true") print("JSON endpoint: curl http://localhost:8000/tasks/list") uvicorn.run(app, host="127.0.0.1", port=8000)