fix: don't hold a database connection for the lifetime of an SSE stream (#28183)

With database session sharing enabled, which the docs recommend for PostgreSQL and for multi-replica deployments, the knowledge pending-files and file process-status endpoints each pinned one pooled connection for as long as their SSE stream stayed open, up to one and two hours respectively. A file wedged in processing keeps a stream open for the full duration, so a handful of users sitting on that page can consume every connection in the pool, and the held transactions sit idle and block autovacuum on those tables.

Both handlers took a request-scoped session for their access checks, and FastAPI only releases a yield dependency once the response body has finished streaming, so the session outlived the handler by the whole life of the stream. Neither generator ever used it. They no longer take that dependency, and the queries they run already open their own short-lived sessions when none is passed. This is the approach the chat completion endpoints already use for the same long-response problem.

Measured against a pool with capacity 11: before, at most 11 concurrent streams could ever be open and every further attempt failed, deterministically across repeat runs. After, 25 of 25 opened. Non-stream latency is unchanged, within run-to-run noise, and behaviour is identical whether session sharing is on or off.
This commit is contained in:
Classic298
2026-08-17 01:46:54 -06:00
committed by GitHub
parent e968445812
commit ba0c4b3932
2 changed files with 13 additions and 11 deletions
+7 -7
View File
@@ -613,9 +613,12 @@ async def get_file_process_status(
id: str,
stream: bool = Query(False),
user=Depends(get_verified_user),
db: AsyncSession = Depends(get_async_session),
):
file = await Files.get_file_by_id(id, db=db)
# NOTE: We intentionally do NOT use Depends(get_async_session) here.
# Database operations manage their own short-lived sessions internally.
# Holding a session here would keep a connection for the entire stream
# (up to two hours) and exhaust the connection pool under concurrent load.
file = await Files.get_file_by_id(id)
if not file:
raise HTTPException(
@@ -623,16 +626,13 @@ async def get_file_process_status(
detail=ERROR_MESSAGES.NOT_FOUND,
)
if file.user_id == user.id or user.role == 'admin' or await has_access_to_file(id, 'read', user, db=db):
if file.user_id == user.id or user.role == 'admin' or await has_access_to_file(id, 'read', user):
if stream:
MAX_FILE_PROCESSING_DURATION = 3600 * 2
async def event_stream(file_id):
# NOTE: We intentionally do NOT capture the request's db session here.
# Each poll creates its own short-lived session to avoid holding a
# connection for hours. A WebSocket push would be more efficient.
for _ in range(MAX_FILE_PROCESSING_DURATION):
file_item = await Files.get_file_by_id(file_id) # Creates own session
file_item = await Files.get_file_by_id(file_id)
if file_item:
data = file_item.model_dump().get('data', {})
status = data.get('status')
+6 -4
View File
@@ -1248,7 +1248,6 @@ async def get_pending_knowledge_files(
id: str,
stream: bool = Query(False),
user=Depends(get_verified_user),
db: AsyncSession = Depends(get_async_session),
):
"""Return files that are being processed for this knowledge base but not yet linked.
@@ -1261,7 +1260,11 @@ async def get_pending_knowledge_files(
When ``stream=true``, returns an SSE stream that polls every 3 seconds
and emits the current pending file list. Closes when no files remain.
"""
knowledge = await Knowledges.get_knowledge_by_id(id=id, db=db)
# NOTE: We intentionally do NOT use Depends(get_async_session) here.
# Database operations manage their own short-lived sessions internally.
# Holding a session here would keep a connection for the entire stream
# (up to an hour) and exhaust the connection pool under concurrent load.
knowledge = await Knowledges.get_knowledge_by_id(id=id)
if not knowledge:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
@@ -1276,7 +1279,6 @@ async def get_pending_knowledge_files(
resource_type='knowledge',
resource_id=knowledge.id,
permission='read',
db=db,
)
):
raise HTTPException(
@@ -1285,7 +1287,7 @@ async def get_pending_knowledge_files(
)
if not stream:
return await Files.get_pending_files_for_knowledge(id, db=db)
return await Files.get_pending_files_for_knowledge(id)
async def event_stream(knowledge_id: str):
MAX_POLL_DURATION = 3600 # 1 hour max