Issue #906: Add timeouts to QP and Translator clients (#1010)

Adds the configs:

- `QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS`
- `TRANSLATOR_CLIENT_TIMEOUT_SECONDS`

With a default of `30`.

Author: bjagg

cloudformation/lif-learner-data-export-api-taskdef-includes.yml

@@ -14,6 +14,10 @@ Environment:
   - Name: LIF_TRANSLATOR_BASE_URL
     Value:
       Fn::Sub: "http://translator-org1.lif.${EnvironmentName}.aws:8007"
+  - Name: QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS
+    Value: "30"
+  - Name: TRANSLATOR_CLIENT_TIMEOUT_SECONDS
+    Value: "30"
 Secrets:
   - Name: LIF_MDR_API_AUTH_TOKEN
     ValueFrom:

components/lif/query_planner_client/core.py

@@ -1,9 +1,30 @@
+import os
+from typing import AsyncGenerator
+
 import httpx
 from lif.exceptions.core import LIFException
 from lif.logging import get_logger
 
 logger = get_logger(__name__)
 
+# Default timeout for Query Planner client API calls (in seconds)
+DEFAULT_QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS = 30
+
+
+def _get_query_planner_timeout_seconds() -> int:
+    return int(os.getenv("QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS", str(DEFAULT_QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS)))
+
+
+async def _get_query_planner_client() -> AsyncGenerator[httpx.AsyncClient]:
+    """
+    Generator that yields an httpx AsyncClient.
+
+    Allows a test harness to override this method to connect to an in-memory Query Planner instance.
+    """
+    timeout = _get_query_planner_timeout_seconds()
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        yield client
+
 
 async def fetch_query_from_query_planner(base_url: str, query: dict) -> list[dict]:
     """POST a LIF query to the Query Planner and return the results.
@@ -14,11 +35,13 @@ async def fetch_query_from_query_planner(base_url: str, query: dict) -> list[dic
     # FUTURE WORK: While the CONFIG has a query_planner_query property,
     # It should be folded into this method (or at least this query_planner_client)
     url = base_url.rstrip("/") + "/query"
+    logger.info("Calling the Query Planner URL: %s with a timeout of %s", url, _get_query_planner_timeout_seconds())
+
     try:
-        async with httpx.AsyncClient() as client:
+        async for client in _get_query_planner_client():
             response = await client.post(url, json=query)
     except httpx.TimeoutException as e:
-        msg = f"Query Planner request timed out: {e}"
+        msg = f"Query Planner request timed out due to: {e}"
         logger.error(msg)
         raise QueryPlannerException(msg) from e
     except httpx.ConnectError as e:

components/lif/translator_client/core.py

@@ -1,10 +1,32 @@
+import os
+from typing import AsyncGenerator
+
 import httpx
 from lif.exceptions.core import LIFException
 from lif.logging import get_logger
 
 logger = get_logger(__name__)
 
 
+# Default timeout for Translator client API calls (in seconds)
+DEFAULT_TRANSLATOR_CLIENT_TIMEOUT_SECONDS = 30
+
+
+def _get_translator_timeout_seconds() -> int:
+    return int(os.getenv("TRANSLATOR_CLIENT_TIMEOUT_SECONDS", str(DEFAULT_TRANSLATOR_CLIENT_TIMEOUT_SECONDS)))
+
+
+async def _get_translator_client() -> AsyncGenerator[httpx.AsyncClient]:
+    """
+    Generator that yields an httpx AsyncClient.
+
+    Allows a test harness to override this method to connect to an in-memory Translator instance.
+    """
+    timeout = _get_translator_timeout_seconds()
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        yield client
+
+
 async def translate_learner_data(
     base_url: str, source_schema_id: str, target_schema_id: str, learner_data: dict, tenant_schema: str | None = None
 ) -> dict:
@@ -15,11 +37,12 @@ async def translate_learner_data(
     """
     url = f"{base_url.rstrip('/')}/translate/source/{source_schema_id}/target/{target_schema_id}"
     headers = {"X-API-Tenant-Schema": tenant_schema} if tenant_schema else {}
+    logger.info("Calling the Translator URL: %s with a timeout of %s", url, _get_translator_timeout_seconds())
     try:
-        async with httpx.AsyncClient() as client:
+        async for client in _get_translator_client():
             response = await client.post(url, json=learner_data, headers=headers)
     except httpx.TimeoutException as e:
-        msg = f"Translator request timed out: {e}"
+        msg = f"Translator request timed out due to: {e}"
         logger.error(msg)
         raise TranslatorException(msg) from e
     except httpx.ConnectError as e:

test/components/lif/query_planner_client/__init__.py

@@ -0,0 +1,108 @@
+import os
+from unittest import mock
+
+import httpx
+import pytest
+from lif.query_planner_client import QueryPlannerException, fetch_query_from_query_planner
+from lif.query_planner_client.core import (
+    DEFAULT_QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS,
+    _get_query_planner_timeout_seconds,
+)
+
+_BASE_URL = "http://localhost:8008"
+_QUERY = {"query": "{ Person { firstName } }"}
+
+
+def _make_http_mock(status_code: int, json_return=None):
+    mock_response = mock.Mock()
+    mock_response.status_code = status_code
+    mock_response.json.return_value = json_return or []
+    mock_response.text = str(json_return or [])
+
+    mock_client = mock.AsyncMock()
+    mock_client.post.return_value = mock_response
+
+    mock_cls = mock.MagicMock()
+    mock_cls.return_value.__aenter__ = mock.AsyncMock(return_value=mock_client)
+    mock_cls.return_value.__aexit__ = mock.AsyncMock(return_value=False)
+    return mock_cls, mock_client
+
+
+async def test_successful_query_returns_json():
+    mock_cls, _ = _make_http_mock(200, [{"firstName": "John"}])
+    with mock.patch("lif.query_planner_client.core.httpx.AsyncClient", mock_cls):
+        result = await fetch_query_from_query_planner(_BASE_URL, _QUERY)
+    assert result == [{"firstName": "John"}]
+
+
+async def test_url_is_constructed_from_base():
+    mock_cls, mock_client = _make_http_mock(200, [])
+    with mock.patch("lif.query_planner_client.core.httpx.AsyncClient", mock_cls):
+        await fetch_query_from_query_planner(_BASE_URL, _QUERY)
+    mock_client.post.assert_called_once()
+    call_url = mock_client.post.call_args[0][0]
+    assert call_url == f"{_BASE_URL}/query"
+
+
+async def test_trailing_slash_stripped_from_base_url():
+    mock_cls, mock_client = _make_http_mock(200, [])
+    with mock.patch("lif.query_planner_client.core.httpx.AsyncClient", mock_cls):
+        await fetch_query_from_query_planner(_BASE_URL + "/", _QUERY)
+    call_url = mock_client.post.call_args[0][0]
+    assert call_url == f"{_BASE_URL}/query"
+
+
+async def test_query_sent_as_json_body():
+    mock_cls, mock_client = _make_http_mock(200, [])
+    with mock.patch("lif.query_planner_client.core.httpx.AsyncClient", mock_cls):
+        await fetch_query_from_query_planner(_BASE_URL, _QUERY)
+    assert mock_client.post.call_args[1]["json"] == _QUERY
+
+
+async def test_non_200_raises_query_planner_exception():
+    mock_cls, _ = _make_http_mock(500)
+    with mock.patch("lif.query_planner_client.core.httpx.AsyncClient", mock_cls):
+        with pytest.raises(QueryPlannerException, match="HTTP 500"):
+            await fetch_query_from_query_planner(_BASE_URL, _QUERY)
+
+
+async def test_timeout_raises_query_planner_exception():
+    with mock.patch("lif.query_planner_client.core.httpx.AsyncClient") as mock_cls:
+        mock_cls.return_value.__aenter__ = mock.AsyncMock(side_effect=httpx.TimeoutException("timed out"))
+        mock_cls.return_value.__aexit__ = mock.AsyncMock(return_value=False)
+        with pytest.raises(QueryPlannerException, match="timed out"):
+            await fetch_query_from_query_planner(_BASE_URL, _QUERY)
+
+
+async def test_connect_error_raises_query_planner_exception():
+    with mock.patch("lif.query_planner_client.core.httpx.AsyncClient") as mock_cls:
+        mock_cls.return_value.__aenter__ = mock.AsyncMock(side_effect=httpx.ConnectError("connection refused"))
+        mock_cls.return_value.__aexit__ = mock.AsyncMock(return_value=False)
+        with pytest.raises(QueryPlannerException, match="connection refused"):
+            await fetch_query_from_query_planner(_BASE_URL, _QUERY)
+
+
+def test_timeout_defaults_when_env_var_absent():
+    with mock.patch.dict(os.environ, {}, clear=True):
+        assert _get_query_planner_timeout_seconds() == DEFAULT_QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS
+
+
+def test_timeout_reads_from_env_var():
+    with mock.patch.dict(os.environ, {"QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS": "5"}):
+        assert _get_query_planner_timeout_seconds() == 5
+
+
+async def test_default_timeout_passed_to_async_client():
+    mock_cls, _ = _make_http_mock(200, [])
+    with mock.patch.dict(os.environ, {}, clear=True):
+        with mock.patch("lif.query_planner_client.core.httpx.AsyncClient", mock_cls):
+            await fetch_query_from_query_planner(_BASE_URL, _QUERY)
+    assert mock_cls.call_args.kwargs["timeout"] == DEFAULT_QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS
+
+
+async def test_configured_timeout_passed_to_async_client():
+    mock_cls, _ = _make_http_mock(200, [])
+    with mock.patch.dict(os.environ, {"QUERY_PLANNER_CLIENT_TIMEOUT_SECONDS": "7"}):
+        with mock.patch("lif.query_planner_client.core.httpx.AsyncClient", mock_cls):
+            await fetch_query_from_query_planner(_BASE_URL, _QUERY)
+    assert mock_cls.call_args.kwargs["timeout"] == 7

test/components/lif/query_planner_client/test_core.py

@@ -1,8 +1,10 @@
+import os
 from unittest import mock
 
 import httpx
 import pytest
 from lif.translator_client import TranslatorException, translate_learner_data
+from lif.translator_client.core import DEFAULT_TRANSLATOR_CLIENT_TIMEOUT_SECONDS, _get_translator_timeout_seconds
 
 _BASE_URL = "http://localhost:8007"
 _SOURCE_ID = "17"
@@ -78,3 +80,29 @@ async def test_connect_error_raises_translator_exception():
         mock_cls.return_value.__aexit__ = mock.AsyncMock(return_value=False)
         with pytest.raises(TranslatorException, match="connection refused"):
             await translate_learner_data(_BASE_URL, _SOURCE_ID, _TARGET_ID, _LEARNER_DATA)
+
+
+def test_timeout_defaults_when_env_var_absent():
+    with mock.patch.dict(os.environ, {}, clear=True):
+        assert _get_translator_timeout_seconds() == DEFAULT_TRANSLATOR_CLIENT_TIMEOUT_SECONDS
+
+
+def test_timeout_reads_from_env_var():
+    with mock.patch.dict(os.environ, {"TRANSLATOR_CLIENT_TIMEOUT_SECONDS": "5"}):
+        assert _get_translator_timeout_seconds() == 5
+
+
+async def test_default_timeout_passed_to_async_client():
+    mock_cls, _ = _make_http_mock(200, {})
+    with mock.patch.dict(os.environ, {}, clear=True):
+        with mock.patch("lif.translator_client.core.httpx.AsyncClient", mock_cls):
+            await translate_learner_data(_BASE_URL, _SOURCE_ID, _TARGET_ID, _LEARNER_DATA)
+    assert mock_cls.call_args.kwargs["timeout"] == DEFAULT_TRANSLATOR_CLIENT_TIMEOUT_SECONDS
+
+
+async def test_configured_timeout_passed_to_async_client():
+    mock_cls, _ = _make_http_mock(200, {})
+    with mock.patch.dict(os.environ, {"TRANSLATOR_CLIENT_TIMEOUT_SECONDS": "7"}):
+        with mock.patch("lif.translator_client.core.httpx.AsyncClient", mock_cls):
+            await translate_learner_data(_BASE_URL, _SOURCE_ID, _TARGET_ID, _LEARNER_DATA)
+    assert mock_cls.call_args.kwargs["timeout"] == 7