diff --git a/shepherd_server/base_routes.py b/shepherd_server/base_routes.py index a424038..e93a015 100644 --- a/shepherd_server/base_routes.py +++ b/shepherd_server/base_routes.py @@ -35,7 +35,11 @@ setup_logging, ) from shepherd_utils.otel import setup_tracer -from shepherd_utils.task_deadline import deadline_field, query_deadline +from shepherd_utils.task_deadline import ( + TIMEOUT_STATUS, + deadline_field, + query_deadline, +) setup_logging() @@ -43,6 +47,23 @@ base_router = APIRouter() +# shepherd_brain.state/status (see shepherd_db/init_db.sql). A query is +# inserted QUEUED/OK and only leaves that state when it finishes, is +# abandoned, or times out. +TERMINAL_QUERY_STATES = {"COMPLETED", "ABANDONED"} +OK_QUERY_STATUS = "OK" +# The status prefix the janitor writes when a query never completed within its +# budget (see the ABANDONED update in shepherd_utils.db). +ABANDONED_STATUS_PREFIX = "abandoned" +# What /query answers with for each way a query can fail. TRAPI 1.5 only +# documents 200/400/429/500/501 for this operation, but a caller is better +# served by the code that actually describes what happened: a query that ran +# out of time is not an internal error, and one that was never accepted because +# the datastore was unavailable is worth retrying. +QUERY_ERROR_CODE = 500 +QUERY_TIMEOUT_CODE = 504 +QUERY_UNAVAILABLE_CODE = 503 + class QueryIntakeError(Exception): """Raised when a query can't be accepted because its initial state could @@ -186,6 +207,41 @@ async def run_query( return query_id, response_id, logger +def query_status_code(status: Optional[str]) -> int: + """The HTTP code describing how a query ended, from its stored status. + + A query that ran out of its budget (``TIMEOUT``) or was reaped without ever + completing (``Abandoned: ...``) is a gateway timeout: Shepherd is fine, the + work behind it didn't finish in time. Anything else non-OK is an operation + that failed, which is a genuine internal error. + """ + if not status or status == OK_QUERY_STATUS: + return 200 + if status == TIMEOUT_STATUS or status.lower().startswith(ABANDONED_STATUS_PREFIX): + return QUERY_TIMEOUT_CODE + return QUERY_ERROR_CODE + + +def apply_query_status(response: dict, status: Optional[str]) -> None: + """Stamp a non-OK query status onto the TRAPI response, in place. + + ``status``/``description`` are TRAPI Response fields, so a caller reading + the body it already parses can tell a failed query from an empty one. An + error the ARA reported itself is left alone -- it is more specific than the + query-level status -- but a body that claims nothing, or claims success for + a query that failed, is corrected here. + """ + if not status or status == OK_QUERY_STATUS: + return + current = response.get("status") + if isinstance(current, str) and ( + "error" in current.lower() or "fail" in current.lower() + ): + return + response["status"] = "Error" + response.setdefault("description", f"Query finished with status {status}.") + + async def run_sync_query( target: ARATargetEnum, query: dict = Body(..., examples=[default_input_query]), @@ -198,7 +254,7 @@ async def run_sync_query( except QueryIntakeError as e: return ORJSONResponse( content={"status": "ERROR", "description": str(e)}, - status_code=500, + status_code=QUERY_UNAVAILABLE_CODE, ) start = time.time() now = start @@ -220,11 +276,24 @@ async def run_sync_query( content={ "status": "ERROR", "description": "Unable to get response", - } + }, + status_code=QUERY_ERROR_CODE, ) logs = await get_logs(response_id, logger) response["logs"] = logs - return ORJSONResponse(content=response) + # The stored status is the one thing that knows the query + # failed -- a response an operation never got to write looks + # exactly like one that legitimately found nothing. Report it + # rather than handing back a body that only says "here you go". + status = query_state[10] + apply_query_status(response, status) + # The body has said "status": "Error" since apply_query_status + # went in, but the HTTP code said 200 -- so a caller that + # checks the code (rather than parsing the payload for a status + # field) saw every failed query as a successful one. + return ORJSONResponse( + content=response, status_code=query_status_code(status) + ) else: # Debug, not warning: this fires every 0.5s while a query is still # in flight (the row just isn't COMPLETED yet) and would otherwise @@ -233,7 +302,10 @@ async def run_sync_query( await asyncio.sleep(0.5) logger.error("Query timed out") - return ORJSONResponse(content={"status": "TIMEOUT", "description": "Query timeout"}) + return ORJSONResponse( + content={"status": "TIMEOUT", "description": "Query timeout"}, + status_code=QUERY_TIMEOUT_CODE, + ) async def run_async_query( @@ -458,14 +530,44 @@ async def callback( @base_router.get("/asyncquery_status/{qid}", status_code=200) async def query_status( qid: str, -) -> dict: +): """Handle query status requests.""" - # TODO: get query status from db - return { - "status": "Queued", - "description": "Query is currently waiting to be run.", - "logs": [], - } + logger = logging.getLogger("shepherd.query_status") + logger.setLevel(logging.INFO) + attach_query_handler(logger) + query_state = await get_query_state(qid, logger) + if query_state is None: + return JSONResponse(content={"error": "Not found"}, status_code=404) + + response_id = query_state[7] + state = query_state[9] + status = query_state[10] + description = query_state[11] + logs = await get_logs(response_id, logger) if response_id else [] + + if state not in TERMINAL_QUERY_STATES: + # Shepherd doesn't track a separate running state: a query is in the + # pipeline from the moment it is accepted until it finishes. + trapi_status = "Running" + default_description = "Query is currently running." + elif status == OK_QUERY_STATUS: + trapi_status = "Completed" + default_description = "Query has finished." + else: + # The query reached the end of the line with something other than OK + # (ERROR from a failed operation, TIMEOUT, ABANDONED). Previously this + # endpoint answered "Queued" for every query it was ever asked about, + # so a failed query and a healthy one looked exactly alike here. + trapi_status = "Failed" + default_description = f"Query finished with status {status}." + + return ORJSONResponse( + content={ + "status": trapi_status, + "description": description or default_description, + "logs": logs, + } + ) @base_router.get("/response/{query_id}", status_code=200) diff --git a/tests/unit/test_query_status.py b/tests/unit/test_query_status.py new file mode 100644 index 0000000..5d53bec --- /dev/null +++ b/tests/unit/test_query_status.py @@ -0,0 +1,233 @@ +"""Tests for how a query's outcome is reported back to the caller. + +``/asyncquery_status/{qid}`` used to answer ``{"status": "Queued"}`` for every +query it was ever asked about (a hardcoded stub with a TODO), so a query that +failed was indistinguishable from a healthy one. The sync path had the milder +version of the same problem: it returned the stored response with no hint that +the query had finished with anything other than OK. +""" + +import json +import logging + +import pytest + +from shepherd_server.base_routes import ( + ARATargetEnum, + QueryIntakeError, + apply_query_status, + query_status, + query_status_code, + run_sync_query, +) + +logger = logging.getLogger(__name__) + + +def _row(state="QUEUED", status="OK", description=None): + """A shepherd_brain row, in the column order get_query_state returns.""" + return ( + "qid", + "start", + "stop", + "submitter", + "ip", + "domain", + "hostname", + "response_id", + None, + state, + status, + description, + ) + + +def _patch_state(mocker, row, logs=None): + mocker.patch( + "shepherd_server.base_routes.get_query_state", + new_callable=mocker.AsyncMock, + return_value=row, + ) + mocker.patch( + "shepherd_server.base_routes.get_logs", + new_callable=mocker.AsyncMock, + return_value=logs if logs is not None else [], + ) + + +def _body(response): + return json.loads(bytes(response.body)) + + +@pytest.mark.asyncio +async def test_status_unknown_query_is_not_found(mocker): + _patch_state(mocker, None) + response = await query_status("qid") + assert response.status_code == 404 + + +@pytest.mark.asyncio +async def test_status_in_flight_query_is_running(mocker): + _patch_state(mocker, _row(state="QUEUED", status="OK")) + assert _body(await query_status("qid"))["status"] == "Running" + + +@pytest.mark.asyncio +async def test_status_finished_query_is_completed(mocker): + _patch_state(mocker, _row(state="COMPLETED", status="OK")) + assert _body(await query_status("qid"))["status"] == "Completed" + + +@pytest.mark.parametrize( + "state,status", + [ + ("COMPLETED", "ERROR"), + ("COMPLETED", "TIMEOUT"), + ("ABANDONED", "Abandoned: no completion within budget"), + ], +) +@pytest.mark.asyncio +async def test_status_failed_query_is_reported_as_failed(mocker, state, status): + _patch_state(mocker, _row(state=state, status=status)) + body = _body(await query_status("qid")) + assert body["status"] == "Failed" + assert status in body["description"] + + +@pytest.mark.asyncio +async def test_status_carries_the_query_logs(mocker): + """The logs are where the upstream status code is recorded.""" + logs = [{"level": "ERROR", "message": "ARAX service returned HTTP 500"}] + _patch_state(mocker, _row(state="COMPLETED", status="ERROR"), logs=logs) + assert _body(await query_status("qid"))["logs"] == logs + + +# --- apply_query_status ---------------------------------------------------- + + +def test_ok_query_is_left_untouched(): + response = {"message": {}} + apply_query_status(response, "OK") + assert response == {"message": {}} + + +def test_failed_query_is_marked_on_the_response(): + response = {"message": {}} + apply_query_status(response, "ERROR") + assert response["status"] == "Error" + assert "ERROR" in response["description"] + + +def test_ara_reported_error_is_not_overwritten(): + """ARAX's own status is more specific than the query-level one.""" + response = {"message": {}, "status": "InternalError", "description": "upstream"} + apply_query_status(response, "ERROR") + assert response["status"] == "InternalError" + assert response["description"] == "upstream" + + +def test_success_claimed_for_a_failed_query_is_corrected(): + response = {"message": {}, "status": "Success"} + apply_query_status(response, "TIMEOUT") + assert response["status"] == "Error" + + +# --- /query ---------------------------------------------------------------- +# +# The body has carried a TRAPI error status since apply_query_status went in, +# but the HTTP code stayed 200, so a caller checking the code rather than +# parsing the payload saw every failed query as a successful one. + + +def _patch_sync_query(mocker, row, response=None): + mocker.patch( + "shepherd_server.base_routes.run_query", + new_callable=mocker.AsyncMock, + return_value=("qid", "response_id", logger), + ) + mocker.patch( + "shepherd_server.base_routes.get_message", + new_callable=mocker.AsyncMock, + return_value=response, + ) + _patch_state(mocker, row) + + +@pytest.mark.asyncio +async def test_query_returns_200_for_a_healthy_query(mocker): + _patch_sync_query( + mocker, _row(state="COMPLETED", status="OK"), response={"message": {}} + ) + response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) + assert response.status_code == 200 + assert "status" not in _body(response) + + +@pytest.mark.parametrize( + "status,code", + [ + # An operation failed: a genuine internal error. + ("ERROR", 500), + # Out of budget, or reaped without ever completing: the work behind + # Shepherd didn't finish in time. + ("TIMEOUT", 504), + ("Abandoned: no completion within budget", 504), + ], +) +@pytest.mark.asyncio +async def test_query_returns_the_code_for_how_it_failed(mocker, status, code): + _patch_sync_query( + mocker, _row(state="COMPLETED", status=status), response={"message": {}} + ) + response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) + assert response.status_code == code + # The body still says which kind of failure it was. + assert _body(response)["status"] == "Error" + assert status in _body(response)["description"] + + +@pytest.mark.asyncio +async def test_query_returns_an_error_code_when_the_response_is_missing(mocker): + _patch_sync_query(mocker, _row(state="COMPLETED", status="OK"), response=None) + response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) + assert response.status_code == 500 + assert _body(response)["description"] == "Unable to get response" + + +@pytest.mark.asyncio +async def test_query_returns_an_error_code_when_it_times_out(mocker): + """The caller's own timeout elapsed with the query still in flight.""" + _patch_sync_query(mocker, _row(state="QUEUED", status="OK")) + response = await run_sync_query( + ARATargetEnum.ARAX, {"message": {}, "parameters": {"timeout": 0}} + ) + assert response.status_code == 504 + assert _body(response)["status"] == "TIMEOUT" + + +@pytest.mark.asyncio +async def test_query_returns_unavailable_when_intake_fails(mocker): + """The query was never accepted, so the caller can retry it as-is.""" + mocker.patch( + "shepherd_server.base_routes.run_query", + new_callable=mocker.AsyncMock, + side_effect=QueryIntakeError("datastore unavailable"), + ) + response = await run_sync_query(ARATargetEnum.ARAX, {"message": {}}) + assert response.status_code == 503 + assert "datastore unavailable" in _body(response)["description"] + + +@pytest.mark.parametrize( + "status,code", + [ + (None, 200), + ("OK", 200), + ("ERROR", 500), + ("TIMEOUT", 504), + ("Abandoned: no completion within budget", 504), + ("something nobody writes today", 500), + ], +) +def test_query_status_code_mapping(status, code): + assert query_status_code(status) == code