diff --git a/sidecar/app/routes/http_helpers.py b/sidecar/app/routes/http_helpers.py index 8d24ba5..bdfe1e4 100644 --- a/sidecar/app/routes/http_helpers.py +++ b/sidecar/app/routes/http_helpers.py @@ -14,3 +14,17 @@ def get_query_from_query_id( if query is None: raise HTTPException(status_code=404, detail=f"Query<{query_id}> not found") return query + + +def check_capacity( + query_manager: QueryManager, +) -> None: + if not query_manager.capacity_available: + raise HTTPException( + status_code=503, + detail=( + f"Capacity unavailable. Currently running " + f"{len(query_manager.running_queries)} of " + f"{query_manager.max_parallel_queries} queries." + ), + ) diff --git a/sidecar/app/routes/start.py b/sidecar/app/routes/start.py index 0d0d99b..25e8784 100644 --- a/sidecar/app/routes/start.py +++ b/sidecar/app/routes/start.py @@ -3,7 +3,7 @@ from pathlib import Path from typing import Annotated -from fastapi import APIRouter, BackgroundTasks, Form, HTTPException, Request, status +from fastapi import APIRouter, BackgroundTasks, Form, Request, status from fastapi.responses import StreamingResponse from ..local_paths import Paths @@ -11,7 +11,7 @@ from ..query.demo_logger import DemoLoggerQuery from ..query.ipa import GateType, IPACoordinatorQuery, IPAHelperQuery from ..settings import get_settings -from .http_helpers import get_query_from_query_id +from .http_helpers import check_capacity, get_query_from_query_id router = APIRouter( prefix="/start", @@ -50,8 +50,8 @@ def demo_logger( request: Request, ): query_manager = request.app.state.QUERY_MANAGER - if not query_manager.capacity_available: - raise HTTPException(status_code=503, detail="Capacity unavailable") + check_capacity(query_manager) + query = DemoLoggerQuery( query_id=query_id, num_lines=num_lines, @@ -74,8 +74,7 @@ def start_ipa_helper( ): # pylint: disable=too-many-arguments query_manager = request.app.state.QUERY_MANAGER - if not query_manager.capacity_available: - raise HTTPException(status_code=503, detail="Capacity unavailable") + check_capacity(query_manager) settings = get_settings() role = settings.role @@ -164,8 +163,7 @@ def start_ipa_query( ): # pylint: disable=too-many-arguments query_manager = request.app.state.QUERY_MANAGER - if not query_manager.capacity_available: - raise HTTPException(status_code=503, detail="Capacity unavailable") + check_capacity(query_manager) settings = get_settings() role = settings.role