From 2fd072ab240179b33372cacc5ce8aa0088272de9 Mon Sep 17 00:00:00 2001 From: Erik Taubeneck Date: Wed, 17 Jul 2024 00:42:13 -0700 Subject: [PATCH] wrapper function for query not found --- sidecar/app/routes/http_helpers.py | 16 ++++++++++++++++ sidecar/app/routes/start.py | 19 ++++++++----------- sidecar/app/routes/stop.py | 13 ++++--------- sidecar/app/routes/websockets.py | 11 +++-------- 4 files changed, 31 insertions(+), 28 deletions(-) create mode 100644 sidecar/app/routes/http_helpers.py diff --git a/sidecar/app/routes/http_helpers.py b/sidecar/app/routes/http_helpers.py new file mode 100644 index 0000000..8d24ba5 --- /dev/null +++ b/sidecar/app/routes/http_helpers.py @@ -0,0 +1,16 @@ +from typing import Type + +from fastapi import HTTPException + +from ..query.base import Query, QueryManager + + +def get_query_from_query_id( + query_manager: QueryManager, + query_cls: Type[Query], + query_id: str, +) -> Query: + query = query_manager.get_from_query_id(query_cls, query_id) + if query is None: + raise HTTPException(status_code=404, detail=f"Query<{query_id}> not found") + return query diff --git a/sidecar/app/routes/start.py b/sidecar/app/routes/start.py index 67c564c..0d0d99b 100644 --- a/sidecar/app/routes/start.py +++ b/sidecar/app/routes/start.py @@ -11,6 +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 router = APIRouter( prefix="/start", @@ -79,7 +80,9 @@ def start_ipa_helper( settings = get_settings() role = settings.role if not role or role == role.COORDINATOR: - raise IncorrectRoleError("Cannot start helper without helper role.") + raise IncorrectRoleError( + f"Cannot start helper without helper role. Currently running {role=}." + ) compiled_id = ( f"{commit_hash}_{gate_type}" @@ -112,10 +115,7 @@ def get_ipa_helper_status( query_id: str, request: Request, ): - query_manager = request.app.state.QUERY_MANAGER - query = query_manager.get_from_query_id(Query, query_id) - if query is None: - raise HTTPException(status_code=404, detail="Query not found") + query = get_query_from_query_id(request.app.state.QUERY_MANAGER, Query, query_id) return {"status": query.status.name} @@ -124,11 +124,7 @@ def get_ipa_helper_log_file( query_id: str, request: Request, ): - query_manager = request.app.state.QUERY_MANAGER - query = query_manager.get_from_query_id(Query, query_id) - if query is None: - raise HTTPException(status_code=404, detail="Query not found") - + query = get_query_from_query_id(request.app.state.QUERY_MANAGER, Query, query_id) settings = get_settings() def iterfile(): @@ -175,7 +171,8 @@ def start_ipa_query( role = settings.role if role != role.COORDINATOR: raise IncorrectRoleError( - f"Sidecar {role}: Cannot start query without coordinator role." + f"Attempting to start query with {role=}: " + "Cannot start query without coordinator role." ) paths = Paths( diff --git a/sidecar/app/routes/stop.py b/sidecar/app/routes/stop.py index 4dd4d54..11ebc1c 100644 --- a/sidecar/app/routes/stop.py +++ b/sidecar/app/routes/stop.py @@ -1,7 +1,8 @@ -from fastapi import APIRouter, HTTPException, Request +from fastapi import APIRouter, Request from ..query.base import Query from ..query.status import Status +from .http_helpers import get_query_from_query_id router = APIRouter( prefix="/stop", @@ -16,10 +17,7 @@ def finish( query_id: str, request: Request, ): - query_manager = request.app.state.QUERY_MANAGER - query = query_manager.get_from_query_id(Query, query_id) - if query is None: - raise HTTPException(status_code=404, detail="Query not found") + query = get_query_from_query_id(request.app.state.QUERY_MANAGER, Query, query_id) query.logger.info(f"{query=}") if query.status < Status.COMPLETE: @@ -34,10 +32,7 @@ def kill( query_id: str, request: Request, ): - query_manager = request.app.state.QUERY_MANAGER - query = query_manager.get_from_query_id(Query, query_id) - if query is None: - raise HTTPException(status_code=404, detail="Query not found") + query = get_query_from_query_id(request.app.state.QUERY_MANAGER, Query, query_id) query.logger.info(f"kill called for {query_id=}") if query.status < Status.COMPLETE: diff --git a/sidecar/app/routes/websockets.py b/sidecar/app/routes/websockets.py index b4307aa..d846841 100644 --- a/sidecar/app/routes/websockets.py +++ b/sidecar/app/routes/websockets.py @@ -6,6 +6,7 @@ from websockets import ConnectionClosedError, ConnectionClosedOK from ..query.base import Query +from .http_helpers import get_query_from_query_id router = APIRouter( prefix="/ws", @@ -51,10 +52,7 @@ async def logs_websocket( websocket: WebSocket, query_id: str, ): - query_manager = websocket.app.state.QUERY_MANAGER - query = query_manager.get_from_query_id(Query, query_id) - if query is None: - raise HTTPException(status_code=404, detail="Query not found") + query = get_query_from_query_id(websocket.app.state.QUERY_MANAGER, Query, query_id) async with use_websocket(websocket) as websocket: with open(query.log_file_path, "r", encoding="utf8") as log_file: @@ -79,10 +77,7 @@ async def stats_websocket( websocket: WebSocket, query_id: str, ): - query_manager = websocket.app.state.QUERY_MANAGER - query = query_manager.get_from_query_id(Query, query_id) - if query is None: - raise HTTPException(status_code=404, detail="Query not found") + query = get_query_from_query_id(websocket.app.state.QUERY_MANAGER, Query, query_id) async with use_websocket(websocket) as websocket: if query.finished: