Propagate ARAX's status code instead of dropping it

tests/unit/test_arax.py

@@ -0,0 +1,218 @@
+"""Tests for ``workers.arax.worker``.
+
+Focused on what happens to the status code ARAX answers with: it used to be
+logged and then dropped -- every failure became the same ``{"status": "error"}``
+blob that the workflow reported as a success -- so nothing downstream had a
+status code to report.
+"""
+
+import json
+import logging
+
+import httpx
+import pytest
+
+from workers.arax.worker import (
+    BAD_GATEWAY,
+    GATEWAY_TIMEOUT,
+    ARAXServiceError,
+    arax,
+)
+
+logger = logging.getLogger(__name__)
+
+QUERY = {
+    "message": {
+        "query_graph": {
+            "nodes": {"a": {"ids": ["MONDO:0005148"]}, "b": {}},
+            "edges": {"e0": {"subject": "a", "object": "b"}},
+        }
+    }
+}
+
+PATHFINDER_QUERY = {
+    "message": {
+        "query_graph": {
+            "nodes": {"a": {"ids": ["MONDO:0005148"]}, "b": {"ids": ["CHEBI:15365"]}},
+            "paths": {"p0": {"subject": "a", "object": "b"}},
+        }
+    }
+}
+
+ARAX_RESPONSE = {
+    "message": {
+        "query_graph": QUERY["message"]["query_graph"],
+        "knowledge_graph": {
+            "nodes": {"MONDO:0005148": {}},
+            "edges": {"e0": {"subject": "a", "object": "b"}},
+        },
+        "results": [],
+    }
+}
+
+
+def _task():
+    return [
+        "task_id",
+        {
+            "query_id": "query_id",
+            "response_id": "response_id",
+            "workflow": json.dumps([{"id": "arax"}]),
+            "log_level": "20",
+            "otel": json.dumps({}),
+            "metadata": json.dumps({}),
+        },
+    ]
+
+
+def _patch_db(mocker, message=None):
+    """Stub the two db calls the worker makes, returning the save mock."""
+    mocker.patch(
+        "workers.arax.worker.get_message",
+        new_callable=mocker.AsyncMock,
+        return_value=message if message is not None else dict(QUERY),
+    )
+    return mocker.patch(
+        "workers.arax.worker.save_message",
+        new_callable=mocker.AsyncMock,
+    )
+
+
+def _patch_post(mocker, response=None, side_effect=None):
+    return mocker.patch(
+        "httpx.AsyncClient.post",
+        new_callable=mocker.AsyncMock,
+        return_value=response,
+        side_effect=side_effect,
+    )
+
+
+def _patch_span(mocker):
+    span = mocker.MagicMock()
+    mocker.patch("workers.arax.worker.get_current_span", return_value=span)
+    return span
+
+
+def _http_response(status_code, json_body=None, text=None):
+    """A real httpx.Response, so is_success/.json()/.content behave as in prod."""
+    request = httpx.Request("POST", "https://arax.example/query")
+    if json_body is not None:
+        return httpx.Response(status_code, json=json_body, request=request)
+    return httpx.Response(status_code, text=text or "", request=request)
+
+
+@pytest.mark.asyncio
+async def test_successful_query_saves_response_and_advances_workflow(mocker):
+    save = _patch_db(mocker)
+    span = _patch_span(mocker)
+    _patch_post(mocker, _http_response(200, json_body=ARAX_RESPONSE))
+    task = _task()
+
+    await arax(task, logger)
+
+    span.set_attribute.assert_any_call("arax.status_code", 200)
+    saved = save.await_args.args[1]
+    assert saved["message"]["results"] == []
+    # Provenance is still injected on the way through.
+    assert saved["message"]["knowledge_graph"]["edges"]["e0"]["sources"] == [
+        {
+            "resource_id": "infores:shepherd-arax",
+            "resource_role": "aggregator_knowledge_source",
+            "source_record_urls": None,
+            "upstream_resource_ids": ["infores:arax"],
+        }
+    ]
+    assert json.loads(task[1]["workflow"]) == [{"id": "arax"}]
+
+
+@pytest.mark.parametrize("status_code", [400, 404, 429, 500, 502])
+@pytest.mark.asyncio
+async def test_error_status_is_propagated(mocker, status_code):
+    """ARAX's own status code reaches the exception, the span and the response."""
+    save = _patch_db(mocker)
+    span = _patch_span(mocker)
+    _patch_post(mocker, _http_response(status_code, text="upstream said no"))
+
+    with pytest.raises(ARAXServiceError) as excinfo:
+        await arax(_task(), logger)
+
+    assert excinfo.value.status_code == status_code
+    assert f"HTTP {status_code}" in str(excinfo.value)
+    assert "upstream said no" in str(excinfo.value)
+    span.set_attribute.assert_any_call("arax.status_code", status_code)
+
+    saved = save.await_args.args[1]
+    assert saved["status"] == "Error"
+    assert f"[HTTP {status_code}]" in saved["description"]
+    # Still a TRAPI response for the query that was asked.
+    assert saved["message"]["query_graph"] == QUERY["message"]["query_graph"]
+    assert saved["message"]["results"] == []
+
+
+@pytest.mark.asyncio
+async def test_error_body_is_truncated(mocker):
+    _patch_db(mocker)
+    _patch_span(mocker)
+    _patch_post(mocker, _http_response(500, text="x" * 5000))
+
+    with pytest.raises(ARAXServiceError) as excinfo:
+        await arax(_task(), logger)
+
+    message = str(excinfo.value)
+    assert "x" * 500 in message and "x" * 501 not in message
+
+
+@pytest.mark.asyncio
+async def test_timeout_reports_gateway_timeout(mocker):
+    save = _patch_db(mocker)
+    span = _patch_span(mocker)
+    _patch_post(mocker, side_effect=httpx.ReadTimeout("timed out"))
+
+    with pytest.raises(ARAXServiceError) as excinfo:
+        await arax(_task(), logger)
+
+    assert excinfo.value.status_code == GATEWAY_TIMEOUT
+    span.set_attribute.assert_any_call("arax.status_code", GATEWAY_TIMEOUT)
+    assert f"[HTTP {GATEWAY_TIMEOUT}]" in save.await_args.args[1]["description"]
+
+
+@pytest.mark.asyncio
+async def test_transport_error_reports_bad_gateway(mocker):
+    save = _patch_db(mocker)
+    span = _patch_span(mocker)
+    _patch_post(mocker, side_effect=httpx.ConnectError(""))
+
+    with pytest.raises(ARAXServiceError) as excinfo:
+        await arax(_task(), logger)
+
+    assert excinfo.value.status_code == BAD_GATEWAY
+    # httpx.ConnectError stringifies to nothing, so the class name carries it.
+    assert "ConnectError" in str(excinfo.value)
+    span.set_attribute.assert_any_call("arax.status_code", BAD_GATEWAY)
+    assert f"[HTTP {BAD_GATEWAY}]" in save.await_args.args[1]["description"]
+
+
+@pytest.mark.asyncio
+async def test_unparseable_success_body_keeps_the_status_code(mocker):
+    save = _patch_db(mocker)
+    _patch_span(mocker)
+    _patch_post(mocker, _http_response(200, text="<html>not json</html>"))
+
+    with pytest.raises(ARAXServiceError) as excinfo:
+        await arax(_task(), logger)
+
+    assert excinfo.value.status_code == 200
+    assert "[HTTP 200]" in save.await_args.args[1]["description"]
+
+
+@pytest.mark.asyncio
+async def test_pathfinder_query_is_routed_without_calling_arax(mocker):
+    save = _patch_db(mocker, message=dict(PATHFINDER_QUERY))
+    post = _patch_post(mocker, _http_response(200, json_body=ARAX_RESPONSE))
+    task = _task()
+
+    await arax(task, logger)
+
+    assert not post.called
+    assert not save.called
+    assert json.loads(task[1]["workflow"]) == [{"id": "arax.pathfinder"}]

workers/arax/worker.py

@@ -5,6 +5,7 @@
 import logging
 import uuid
 import httpx
+from opentelemetry.trace import get_current_span
 from shepherd_utils.inject_shepherd_arax_provenance import (
     add_shepherd_arax_to_edge_sources,
 )
@@ -20,8 +21,47 @@
 GROUP = "consumer"
 CONSUMER = str(uuid.uuid4())[:8]
 TASK_LIMIT = 10
+ARAX_TIMEOUT = 270
 tracer = setup_tracer(STREAM)
 LOGGER = get_worker_logger(STREAM)
+# How much of a failing ARAX response body goes into the error. The body of a
+# non-2xx can be an arbitrarily large HTML error page, and this string ends up
+# in the query's logs, so keep only the head of it.
+ERROR_BODY_BYTES = 500
+# Used when ARAX never answered at all, so there is no status code of its own to
+# pass on: we are a gateway in front of it, and these are the codes that say so.
+BAD_GATEWAY = 502
+GATEWAY_TIMEOUT = 504
+
+
+class ARAXServiceError(Exception):
+    """The ARAX service did not return a usable TRAPI response.
+
+    Carries the upstream HTTP status code so it survives the whole way out:
+    onto the task span as ``arax.status_code``, into the query's logs via
+    ``run_task_lifecycle``, and so into what the caller gets back. Previously
+    the status code was only logged and the failure was swallowed into a
+    non-TRAPI ``{"status": "error"}`` blob that the workflow then reported as a
+    successful response, which left nothing downstream with a status code to
+    report (see the same fix in ``arax_pathfinder``).
+    """
+
+    def __init__(self, message: str, status_code: int):
+        super().__init__(message)
+        self.status_code = status_code
+
+
+def body_head(response: httpx.Response) -> str:
+    """The first ``ERROR_BODY_BYTES`` of a response body, for an error message."""
+    try:
+        body = response.content[:ERROR_BODY_BYTES]
+    except Exception:
+        # Body not readable (streamed/closed response) -- the status code is
+        # still worth reporting on its own.
+        return ""
+    if not body:
+        return ""
+    return f": {body.decode('utf-8', 'replace')}"
 
 
 def is_pathfinder_query(message):
@@ -44,28 +84,105 @@ def is_pathfinder_query(message):
     return len(qpaths) == 1
 
 
+async def call_arax(message: dict, logger: logging.Logger) -> dict:
+    """POST the message to the ARAX service and return its TRAPI response.
+
+    Raises ``ARAXServiceError`` -- carrying ARAX's own status code -- for
+    anything that isn't a parseable 2xx, so the code reaches the span, the
+    query's logs and the response instead of being logged and dropped.
+    """
+    message["submitter"] = "Shepherd"
+    logger.info(f"Get the message from db {message}")
+    headers = {"Content-Type": "application/json"}
+    span = get_current_span()
+    try:
+        async with httpx.AsyncClient(timeout=ARAX_TIMEOUT) as client:
+            response = await client.post(
+                settings.arax_url, json=message, headers=headers
+            )
+    except httpx.TimeoutException as e:
+        span.set_attribute("arax.status_code", GATEWAY_TIMEOUT)
+        raise ARAXServiceError(
+            f"ARAX service at {settings.arax_url} did not respond within "
+            f"{ARAX_TIMEOUT}s: {type(e).__name__}",
+            GATEWAY_TIMEOUT,
+        ) from e
+    except Exception as e:
+        # httpx reports connect failures, TLS errors and protocol errors as
+        # distinct classes, several of which stringify to an empty message --
+        # hence the type name alongside the message.
+        span.set_attribute("arax.status_code", BAD_GATEWAY)
+        raise ARAXServiceError(
+            f"Error occurred calling ARAX service at {settings.arax_url}: "
+            f"{type(e).__name__}: {e}",
+            BAD_GATEWAY,
+        ) from e
+
+    status_code = response.status_code
+    span.set_attribute("arax.status_code", status_code)
+    logger.info(f"Status Code from ARAX response: {status_code}")
+    if not response.is_success:
+        raise ARAXServiceError(
+            f"ARAX service at {settings.arax_url} returned HTTP "
+            f"{status_code}{body_head(response)}",
+            status_code,
+        )
+
+    try:
+        result = response.json()
+    except Exception as e:
+        # A 2xx whose body isn't TRAPI JSON is still a failed lookup, and
+        # ARAX's status code is the most useful thing we know about it.
+        raise ARAXServiceError(
+            f"ARAX service at {settings.arax_url} returned HTTP {status_code} "
+            f"with a body that could not be parsed as JSON: {e}",
+            status_code,
+        ) from e
+
+    return add_shepherd_arax_to_edge_sources(result)
+
+
+def error_response(message: dict, error: ARAXServiceError) -> dict:
+    """A TRAPI response reporting a failed ARAX call.
+
+    ``status``/``description`` are TRAPI Response fields, so the status code
+    lands somewhere the caller already parses rather than only in the logs. The
+    query graph is carried over and the result containers are emptied, so what
+    comes back is still a valid TRAPI response for the query that was asked.
+    """
+    query_graph = {}
+    if isinstance(message.get("message"), dict):
+        query_graph = message["message"].get("query_graph") or {}
+    return {
+        "message": {
+            "query_graph": query_graph,
+            "knowledge_graph": {"nodes": {}, "edges": {}},
+            "results": [],
+        },
+        "status": "Error",
+        "description": f"[HTTP {error.status_code}] {error}",
+    }
+
+
 async def arax(task, logger: logging.Logger):
     query_id = task[1]["query_id"]
     logger.info(f"Getting message from db for query id {query_id}")
     message = await get_message(query_id, logger)
     if is_pathfinder_query(message):
         task[1]["workflow"] = json.dumps([{"id": "arax.pathfinder"}])
     else:
-        try:
-            message["submitter"] = "Shepherd"
-            logger.info(f"Get the message from db {message}")
-            headers = {"Content-Type": "application/json"}
-            async with httpx.AsyncClient(timeout=270) as client:
-                response = await client.post(
-                    settings.arax_url, json=message, headers=headers
-                )
-            logger.info(f"Status Code from ARAX response: {response.status_code}")
-            result = response.json()
-            result = add_shepherd_arax_to_edge_sources(result)
-        except Exception as e:
-            logger.error(f"Error occurred calling ARAX service: {e}")
-            result = {"status": "error", "error": str(e)}
         response_id = task[1]["response_id"]
+        try:
+            result = await call_arax(message, logger)
+        except ARAXServiceError as e:
+            # Leave the caller a TRAPI response that says what happened before
+            # letting the failure reach run_task_lifecycle, which records it on
+            # the span and routes the query to finish_query with an ERROR
+            # status. Without this the response id still holds the echo of the
+            # incoming query, so the caller gets a query that looks like it
+            # simply found nothing.
+            await save_message(response_id, error_response(message, e), logger)
+            raise
         await save_message(response_id, result, logger)
         task[1]["workflow"] = json.dumps([{"id": "arax"}])