From 503eb8f496a1c2f08600759c3d13f8787a6689c0 Mon Sep 17 00:00:00 2001 From: Evan Morris Date: Tue, 1 Sep 2026 10:46:08 -0400 Subject: [PATCH 1/2] use orjson for query endpoints --- shepherd_server/aras/aragorn.py | 20 ++-- shepherd_server/aras/arax.py | 20 ++-- shepherd_server/aras/bte.py | 20 ++-- shepherd_server/aras/sipr.py | 20 ++-- shepherd_server/base_routes.py | 93 ++++++++++++++- tests/unit/test_query_body_parsing.py | 163 ++++++++++++++++++++++++++ 6 files changed, 283 insertions(+), 53 deletions(-) create mode 100644 tests/unit/test_query_body_parsing.py diff --git a/shepherd_server/aras/aragorn.py b/shepherd_server/aras/aragorn.py index b8a96cd..2dd2541 100644 --- a/shepherd_server/aras/aragorn.py +++ b/shepherd_server/aras/aragorn.py @@ -1,4 +1,4 @@ -from fastapi import Body, FastAPI, Request, Response +from fastapi import FastAPI, Request, Response from fastapi.openapi.docs import ( get_swagger_ui_html, ) @@ -7,7 +7,7 @@ from shepherd_server.base_routes import ( ARATargetEnum, base_router, - default_input_query, + query_openapi_extra, run_async_query, run_sync_query, callback, @@ -17,19 +17,15 @@ ARAGORN = FastAPI(title="Shepherd Aragorn") -@ARAGORN.post("/query") -async def sync_query( - query: dict = Body(..., examples=[default_input_query]), -) -> Response: - response = await run_sync_query(ARATargetEnum.ARAGORN, query) +@ARAGORN.post("/query", openapi_extra=query_openapi_extra) +async def sync_query(request: Request) -> Response: + response = await run_sync_query(ARATargetEnum.ARAGORN, request) return response -@ARAGORN.post("/asyncquery") -async def async_query( - query: dict = Body(..., examples=[default_input_query]), -) -> Response: - response = await run_async_query(ARATargetEnum.ARAGORN, query) +@ARAGORN.post("/asyncquery", openapi_extra=query_openapi_extra) +async def async_query(request: Request) -> Response: + response = await run_async_query(ARATargetEnum.ARAGORN, request) return response diff --git a/shepherd_server/aras/arax.py b/shepherd_server/aras/arax.py index ceaa2d0..52f1421 100644 --- a/shepherd_server/aras/arax.py +++ b/shepherd_server/aras/arax.py @@ -1,4 +1,4 @@ -from fastapi import Body, FastAPI, Request, Response +from fastapi import FastAPI, Request, Response from fastapi.openapi.docs import ( get_swagger_ui_html, ) @@ -7,7 +7,7 @@ from shepherd_server.base_routes import ( ARATargetEnum, base_router, - default_input_query, + query_openapi_extra, run_async_query, run_sync_query, ) @@ -16,19 +16,15 @@ ARAX = FastAPI(title="Shepherd ARAX") -@ARAX.post("/query") -async def sync_query( - query: dict = Body(..., examples=[default_input_query]), -) -> Response: - response = await run_sync_query(ARATargetEnum.ARAX, query) +@ARAX.post("/query", openapi_extra=query_openapi_extra) +async def sync_query(request: Request) -> Response: + response = await run_sync_query(ARATargetEnum.ARAX, request) return response -@ARAX.post("/asyncquery") -async def async_query( - query: dict = Body(..., examples=[default_input_query]), -) -> Response: - response = await run_async_query(ARATargetEnum.ARAX, query) +@ARAX.post("/asyncquery", openapi_extra=query_openapi_extra) +async def async_query(request: Request) -> Response: + response = await run_async_query(ARATargetEnum.ARAX, request) return response diff --git a/shepherd_server/aras/bte.py b/shepherd_server/aras/bte.py index 0799e26..46bb76e 100644 --- a/shepherd_server/aras/bte.py +++ b/shepherd_server/aras/bte.py @@ -1,4 +1,4 @@ -from fastapi import Body, FastAPI, Request, Response +from fastapi import FastAPI, Request, Response from fastapi.openapi.docs import ( get_swagger_ui_html, ) @@ -8,7 +8,7 @@ ARATargetEnum, base_router, callback, - default_input_query, + query_openapi_extra, run_async_query, run_sync_query, ) @@ -17,19 +17,15 @@ BTE = FastAPI(title="Shepherd BTE") -@BTE.post("/query") -async def sync_query( - query: dict = Body(..., examples=[default_input_query]), -) -> Response: - response = await run_sync_query(ARATargetEnum.BTE, query) +@BTE.post("/query", openapi_extra=query_openapi_extra) +async def sync_query(request: Request) -> Response: + response = await run_sync_query(ARATargetEnum.BTE, request) return response -@BTE.post("/asyncquery") -async def async_query( - query: dict = Body(..., examples=[default_input_query]), -) -> Response: - response = await run_async_query(ARATargetEnum.BTE, query) +@BTE.post("/asyncquery", openapi_extra=query_openapi_extra) +async def async_query(request: Request) -> Response: + response = await run_async_query(ARATargetEnum.BTE, request) return response diff --git a/shepherd_server/aras/sipr.py b/shepherd_server/aras/sipr.py index 1da8db0..76e9a90 100644 --- a/shepherd_server/aras/sipr.py +++ b/shepherd_server/aras/sipr.py @@ -1,4 +1,4 @@ -from fastapi import Body, FastAPI, Request, Response +from fastapi import FastAPI, Request, Response from fastapi.openapi.docs import ( get_swagger_ui_html, ) @@ -7,7 +7,7 @@ from shepherd_server.base_routes import ( ARATargetEnum, base_router, - default_input_query, + query_openapi_extra, run_async_query, run_sync_query, ) @@ -16,19 +16,15 @@ SIPR = FastAPI(title="Shepherd SIPR") -@SIPR.post("/query") -async def sync_query( - query: dict = Body(..., examples=[default_input_query]), -) -> Response: - response = await run_sync_query(ARATargetEnum.SIPR, query) +@SIPR.post("/query", openapi_extra=query_openapi_extra) +async def sync_query(request: Request) -> Response: + response = await run_sync_query(ARATargetEnum.SIPR, request) return response -@SIPR.post("/asyncquery") -async def async_query( - query: dict = Body(..., examples=[default_input_query]), -) -> Response: - response = await run_async_query(ARATargetEnum.SIPR, query) +@SIPR.post("/asyncquery", openapi_extra=query_openapi_extra) +async def async_query(request: Request) -> Response: + response = await run_async_query(ARATargetEnum.SIPR, request) return response diff --git a/shepherd_server/base_routes.py b/shepherd_server/base_routes.py index 04041a9..ec0af98 100644 --- a/shepherd_server/base_routes.py +++ b/shepherd_server/base_routes.py @@ -10,7 +10,7 @@ import orjson import zstandard -from fastapi import APIRouter, Body, Request, Response +from fastapi import APIRouter, Request, Response from fastapi.responses import JSONResponse, ORJSONResponse from opentelemetry.propagate import extract, inject @@ -52,6 +52,13 @@ class QueryIntakeError(Exception): logged server-side (see ``PG_DISK_FULL`` in ``shepherd_utils.db``).""" +class QueryBodyError(Exception): + """Raised when a posted query body isn't a JSON object. + + Client-safe like ``QueryIntakeError``: the message is returned verbatim as + the 422 ``detail``.""" + + class ARATargetEnum(str, Enum): ARAGORN = "aragorn" ARAX = "arax" @@ -82,6 +89,76 @@ class ARATargetEnum(str, Enum): } +# OpenAPI overrides for /query and /asyncquery. +# +# Those routes take the raw ``Request`` so the body can be parsed with orjson +# (see ``parse_query_body``), which leaves FastAPI with no body parameter to +# infer a schema from -- and so no auto-generated request body or 422. Declaring +# both here, and wiring them in via each route's ``openapi_extra``, keeps +# /openapi.json equivalent to what the old +# ``query: dict = Body(..., examples=[default_input_query])`` signature +# produced, so the TRAPI validators and the Swagger "Try it out" example are +# unaffected by the parser swap. +# +# The 422 body is documented as ``{"detail": str}`` rather than FastAPI's +# ``HTTPValidationError`` (``{"detail": [ValidationError, ...]}``): the rejection +# is ours now, and it carries a single message. +query_openapi_extra: dict = { + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": True, + "type": "object", + "title": "Query", + "examples": [default_input_query], + } + } + }, + "required": True, + }, + "responses": { + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "type": "object", + "title": "QueryValidationError", + "properties": {"detail": {"type": "string"}}, + } + } + }, + } + }, +} + + +async def parse_query_body(request: Request) -> dict: + """Read and parse a posted TRAPI query body. + + Deliberately bypasses FastAPI's own body handling, which routes through + Starlette's ``Request.json()`` -> stdlib ``json.loads``. TRAPI query bodies + routinely carry a populated knowledge graph, and stdlib parsing of one costs + ~100ms per 13MB -- all of it blocking the single event loop this server runs + on, so it delays every other in-flight request too. orjson is ~30% faster on + the same payload, and skipping FastAPI's ``dict`` body field drops a + pydantic pass that only ever re-copied the four top-level keys. + + Raises ``QueryBodyError`` for anything that isn't a JSON object. The old + ``query: dict`` annotation got that check for free from pydantic; without it + a posted list or string would reach ``query.get(...)`` downstream and 500. + """ + raw = await request.body() + try: + query = orjson.loads(raw) + except orjson.JSONDecodeError as e: + raise QueryBodyError("Invalid request body: not valid JSON.") from e + if not isinstance(query, dict): + raise QueryBodyError("Invalid request body: expected a JSON object.") + return query + + async def run_query( target: str, query: dict, @@ -177,11 +254,13 @@ async def run_query( async def run_sync_query( target: ARATargetEnum, - query: dict = Body(..., examples=[default_input_query]), + request: Request, ) -> Response: """Handle synchronous TRAPI queries.""" - # query_dict = query.dict() - query_dict = query + try: + query_dict = await parse_query_body(request) + except QueryBodyError as e: + return ORJSONResponse(content={"detail": str(e)}, status_code=422) try: query_id, response_id, logger = await run_query(target, query_dict) except QueryIntakeError as e: @@ -227,9 +306,13 @@ async def run_sync_query( async def run_async_query( target: ARATargetEnum, - query: dict = Body(..., examples=[default_input_query]), + request: Request, ) -> JSONResponse: """Handle asynchronous TRAPI queries.""" + try: + query = await parse_query_body(request) + except QueryBodyError as e: + return JSONResponse(content={"detail": str(e)}, status_code=422) callback_url = query.get("callback") if callback_url is None: return JSONResponse( diff --git a/tests/unit/test_query_body_parsing.py b/tests/unit/test_query_body_parsing.py new file mode 100644 index 0000000..514a973 --- /dev/null +++ b/tests/unit/test_query_body_parsing.py @@ -0,0 +1,163 @@ +"""Tests for /query and /asyncquery body parsing. + +Both routes take the raw ``Request`` and parse it with orjson rather than +letting FastAPI do it, because Starlette's ``Request.json()`` goes through +stdlib ``json.loads`` -- slow on the populated knowledge graphs TRAPI bodies +carry, and blocking on the single event loop the server runs on. + +Dropping the ``query: dict = Body(...)`` signature also dropped the pydantic +pass that used to reject a non-object body, so these tests pin the replacement +checks in ``parse_query_body`` and confirm the OpenAPI request body the routes +now declare by hand still matches what that signature generated. +""" + +from unittest import mock + +import pytest +from fastapi.testclient import TestClient + +from shepherd_server import base_routes +from shepherd_server.aras.aragorn import ARAGORN +from shepherd_server.aras.arax import ARAX +from shepherd_server.aras.bte import BTE +from shepherd_server.aras.sipr import SIPR +from shepherd_server.base_routes import ( + QueryBodyError, + default_input_query, + parse_query_body, +) + +from .test_callback_size_limit import _make_request + + +ALL_APPS = (ARAGORN, ARAX, BTE, SIPR) + + +# --- parse_query_body ---------------------------------------------------- + + +async def test_parses_a_json_object(): + request = _make_request(b'{"message": {"results": []}}', {}) + assert await parse_query_body(request) == {"message": {"results": []}} + + +@pytest.mark.parametrize("body", [b"{not json", b"", b'{"message":']) +async def test_rejects_malformed_json(body): + with pytest.raises(QueryBodyError, match="not valid JSON"): + await parse_query_body(_make_request(body, {})) + + +@pytest.mark.parametrize("body", [b"[1, 2, 3]", b'"hello"', b"42", b"null"]) +async def test_rejects_non_object_json(body): + # pydantic used to reject these for free via the ``dict`` annotation. + # Without the check, they'd reach ``query.get(...)`` downstream and 500. + with pytest.raises(QueryBodyError, match="expected a JSON object"): + await parse_query_body(_make_request(body, {})) + + +# --- route behavior ------------------------------------------------------ + + +@pytest.fixture +def client_factory(): + """Build a TestClient with query intake stubbed out. + + ``run_query`` is patched at the ``base_routes`` name the route handlers + resolve, so nothing touches Redis or Postgres; the captured call lets the + tests assert the parsed body arrived intact. + """ + captured = {} + + def _factory(app): + async def fake_run_query(target, query, callback_url=None): + captured.update(target=target, query=query, callback_url=callback_url) + return "qid12345", "rid12345", mock.MagicMock() + + patch = mock.patch.object(base_routes, "run_query", fake_run_query) + patch.start() + return TestClient(app), captured, patch + + yield _factory + mock.patch.stopall() + + +@pytest.mark.parametrize("app", ALL_APPS) +def test_asyncquery_passes_the_parsed_body_through(app, client_factory): + client, captured, _ = client_factory(app) + body = dict(default_input_query, callback="http://callback/1") + + response = client.post("/asyncquery", json=body) + + assert response.status_code == 200 + assert response.json()["status"] == "Accepted" + assert response.json()["job_id"] == "qid12345" + # orjson round-trips the whole body, nested query graph included. + assert captured["query"] == body + assert captured["callback_url"] == "http://callback/1" + + +@pytest.mark.parametrize("app", ALL_APPS) +def test_asyncquery_still_requires_a_callback(app, client_factory): + """The pre-existing missing-callback 422 is unchanged by the parser swap.""" + client, _, _ = client_factory(app) + + response = client.post("/asyncquery", json=default_input_query) + + assert response.status_code == 422 + assert response.json() == { + "status": "Failed", + "description": "callback URL missing", + } + + +@pytest.mark.parametrize("path", ["/query", "/asyncquery"]) +def test_malformed_body_returns_422(path, client_factory): + client, captured, _ = client_factory(ARAGORN) + + response = client.post( + path, content=b"{not json", headers={"content-type": "application/json"} + ) + + assert response.status_code == 422 + assert "not valid JSON" in response.json()["detail"] + # Rejected before intake, so no query was ever registered. + assert captured == {} + + +@pytest.mark.parametrize("path", ["/query", "/asyncquery"]) +def test_non_object_body_returns_422(path, client_factory): + client, captured, _ = client_factory(ARAGORN) + + response = client.post(path, json=[1, 2, 3]) + + assert response.status_code == 422 + assert "expected a JSON object" in response.json()["detail"] + assert captured == {} + + +# --- OpenAPI ------------------------------------------------------------- + + +@pytest.mark.parametrize("app", ALL_APPS) +@pytest.mark.parametrize("path", ["/query", "/asyncquery"]) +def test_openapi_still_documents_the_request_body(app, path): + """The hand-written ``openapi_extra`` must match what ``Body(...)`` produced. + + Taking a raw ``Request`` leaves FastAPI nothing to infer a schema from, so + an unnoticed regression here would silently publish a TRAPI endpoint with no + documented request body. + """ + operation = app.openapi()["paths"][path]["post"] + + request_body = operation["requestBody"] + assert request_body["required"] is True + schema = request_body["content"]["application/json"]["schema"] + assert schema["type"] == "object" + assert schema["additionalProperties"] is True + assert schema["examples"] == [default_input_query] + + assert "422" in operation["responses"] + error_schema = operation["responses"]["422"]["content"]["application/json"][ + "schema" + ] + assert error_schema["properties"] == {"detail": {"type": "string"}} \ No newline at end of file From 3bda0505068e858ec932505218a3bf6f94e3dfa8 Mon Sep 17 00:00:00 2001 From: Evan Morris Date: Fri, 4 Sep 2026 17:38:53 -0400 Subject: [PATCH 2/2] eof newline --- tests/unit/test_query_body_parsing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_query_body_parsing.py b/tests/unit/test_query_body_parsing.py index 514a973..504e24c 100644 --- a/tests/unit/test_query_body_parsing.py +++ b/tests/unit/test_query_body_parsing.py @@ -160,4 +160,4 @@ def test_openapi_still_documents_the_request_body(app, path): error_schema = operation["responses"]["422"]["content"]["application/json"][ "schema" ] - assert error_schema["properties"] == {"detail": {"type": "string"}} \ No newline at end of file + assert error_schema["properties"] == {"detail": {"type": "string"}}