Refactor PseudoClient to use async for SID endpoint and update valida… (#480)

* Refactor PseudoClient to use async for SID endpoint and update validation logic

* Add arguments to docstrings

* bump version

---------

Co-authored-by: Nicholas <3789764+skykanin@users.noreply.github.com>

Author: ArielKarlsen

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dapla-toolbelt-pseudo"
-version = "6.0.1"
+version = "6.0.2"
 description = "Pseudonymization extensions for Dapla"
 authors = [{ name = "Dapla Developers", email = "dapla-platform-developers@ssb.no" }]
 requires-python = ">=3.11,<4.0"

src/dapla_pseudo/v1/client.py

@@ -287,29 +287,77 @@ def _handle_response_error_sync(response: requests.Response) -> None:
                 print(response.text)
                 response.raise_for_status()
 
-    def _post_to_sid_endpoint(
+    async def _post_to_sid_endpoint(
         self,
         path: str,
         values: list[str],
         sid_snapshot_date: date | None = None,
-        stream: bool = True,
-    ) -> requests.Response:
-        request: dict[str, t.Collection[str]] = {"fnrList": values}
-        response = requests.post(
-            url=f"{self.pseudo_service_url}/{path}",
-            params={"snapshot": str(sid_snapshot_date)} if sid_snapshot_date else None,
-            # Do not set content-type, as this will cause the json to serialize incorrectly
-            headers={
-                "Authorization": f"Bearer {self.__auth_token()}",
-                "X-Correlation-Id": PseudoClient._generate_new_correlation_id(),
-            },
-            json=request,
-            stream=stream,
-            timeout=TIMEOUT_DEFAULT,  # seconds
-        )
+    ) -> tuple[list[str], str | None]:
+        """Post SID lookup in batches concurrently and merge responses.
+
+        Args:
+            path: Endpoint path to append to the base pseudo_service_url
+            values: List of FNR to look up in the SID-catalogue.
+            sid_snapshot_date: Date representing SID-catalogue version to use. Latest if unspecified. Format: YYYY-MM-DD
+
+        Returns:
+            tuple[list[str], str | None]: (missing_values, datasetExtractionSnapshotTime)
+        """
+        total_rows = len(values)
+        batch_size = self.rows_per_partition
+
+        # Do not split if total rows is less than batch size (default 10k)
+        if total_rows <= batch_size:
+            batches = [values]
+        else:
+            batches = [
+                values[i : i + batch_size] for i in range(0, total_rows, batch_size)
+            ]
+
+        async with ClientSession(
+            connector=TCPConnector(limit=100, enable_cleanup_closed=True),
+            timeout=ClientTimeout(total=TIMEOUT_DEFAULT),
+        ) as session:
 
-        PseudoClient._handle_response_error_sync(response)
-        return response
+            async def _post_batch(
+                batch: list[str],
+                path: str,
+            ) -> dict[str, list[str] | str]:
+                resp_cm = await session.post(
+                    url=f"{self.pseudo_service_url}/{path}",
+                    params=(
+                        {"snapshot": str(sid_snapshot_date)}
+                        if sid_snapshot_date
+                        else None
+                    ),
+                    # Do not set content-type, as this will cause the json to serialize incorrectly
+                    headers={
+                        "Authorization": f"Bearer {self.__auth_token()}",
+                        "X-Correlation-Id": PseudoClient._generate_new_correlation_id(),
+                    },
+                    json={"fnrList": batch},
+                )
+                async with resp_cm as response:
+                    await PseudoClient._handle_response_error(response)
+                    payload = await response.json()
+                    return t.cast(
+                        dict[str, list[str] | str],
+                        (
+                            payload[0]
+                            if isinstance(payload, list) and payload
+                            else payload
+                        ),
+                    )
+
+            results = await asyncio.gather(*[_post_batch(b, path) for b in batches])
+
+        # Merge results
+        all_missing = [m for r in results for m in r.get("missing", [])]
+        raw_snapshot = (
+            results[0].get("datasetExtractionSnapshotTime") if results else None
+        )
+        snapshot_time = t.cast(str | None, raw_snapshot)
+        return all_missing, snapshot_time
 
 
 def _client() -> PseudoClient:

src/dapla_pseudo/v1/validation.py

@@ -1,14 +1,12 @@
 """Builder for submitting a validation request."""
 
-import json
-from collections.abc import Sequence
+import asyncio
 from datetime import date
 from pathlib import Path
 from typing import Any
 
 import pandas as pd
 import polars as pl
-import requests
 
 from dapla_pseudo.utils import convert_to_date
 from dapla_pseudo.utils import get_file_format_from_file_name
@@ -98,7 +96,7 @@ def __init__(
             self._dataframe: pl.DataFrame = dataframe
             self._field: str = field
 
-        def validate_map_to_stable_id(
+        async def _validate_map_to_stable_id_async(
             self, sid_snapshot_date: str | date | None = None
         ) -> Result:
             """Checks if all the selected fields can be mapped to a stable ID.
@@ -112,31 +110,26 @@ def validate_map_to_stable_id(
             """
             Validator._ensure_field_valid(self._field, self._dataframe)
 
-            response: requests.Response = _client()._post_to_sid_endpoint(
-                "sid/lookup/batch",
-                self._dataframe[self._field].to_list(),
-                convert_to_date(sid_snapshot_date),
-                stream=True,
+            client = _client()
+            all_values = self._dataframe[self._field].to_list()
+            snapshot = convert_to_date(sid_snapshot_date)
+
+            missing, extraction_time = await client._post_to_sid_endpoint(
+                path="sid/lookup/batch",
+                values=all_values,
+                sid_snapshot_date=snapshot,
+            )
+
+            result_df = pl.Series(self._field, missing).to_frame()
+            metadata_logs: list[str] = (
+                [f"SID snapshot time {extraction_time}"] if extraction_time else []
             )
-            # The response content is received as a buffered byte stream from the server.
-            # We decode the content using UTF-8, which gives us a List[Dict[str]] structure.
-            result_json = json.loads(response.content.decode("utf-8"))[0]
-            result: Sequence[str] = []
-            metadata: list[str] = []
-            if "missing" in result_json:
-                result = result_json["missing"]
-            if "datasetExtractionSnapshotTime" in result_json:
-                extraction_time = result_json["datasetExtractionSnapshotTime"]
-                metadata = [f"SID snapshot time {extraction_time}"]
-
-            result_df = pl.Series(self._field, result).to_frame()
-            # TODO - make the Validator fit the Result() class better
             return Result(
                 PseudoFieldResponse(
                     data=result_df,
                     raw_metadata=[
                         RawPseudoMetadata(
-                            logs=metadata,
+                            logs=metadata_logs,
                             metrics=[],
                             datadoc=[],
                             field_name=self._field,
@@ -145,6 +138,12 @@ def validate_map_to_stable_id(
                 )
             )
 
+        def validate_map_to_stable_id(
+            self, sid_snapshot_date: str | date | None = None
+        ) -> Result:
+            """Validate mapping to SID (sync wrapper around the async batched version)."""
+            return asyncio.run(self._validate_map_to_stable_id_async(sid_snapshot_date))
+
     @staticmethod
     def _ensure_field_valid(field: str, dataframe: pl.DataFrame) -> None:
         """Ensure that all values are numeric and valid.

tests/v1/integration/test_large_file.py

@@ -0,0 +1,32 @@
+import polars as pl
+import pytest
+import requests
+from tests.v1.integration.utils import integration_test
+
+from dapla_pseudo import Validator
+from dapla_pseudo.v1.result import Result
+
+
+@pytest.mark.usefixtures("setup")
+@integration_test()
+def test_sid_lookup_batch_payload_too_large() -> None:
+    n_rows = 1_400_000
+    df = pl.DataFrame({"fnr": [f"{i:011d}" for i in range(n_rows)]})
+
+    try:
+        result = Validator.from_polars(df).on_field("fnr").validate_map_to_stable_id()
+    except requests.HTTPError as err:
+        # Explicitly fail on HTTP 413
+        if err.response is not None and err.response.status_code == 413:
+            pytest.fail(f"Unexpected 413 Payload Too Large: {err.response.text}")
+        # Re-raise other HTTP errors to surface genuine failures
+        raise
+    except requests.ConnectionError as err:
+        msg = str(err)
+        # Some servers may close connection with a payload-too-large reason
+        if ("Payload Too Large" in msg) or ("Request Entity Too Large" in msg):
+            pytest.fail(f"Unexpected payload-too-large connection error: {msg}")
+        raise
+
+    # Sanity check: we got a Result back
+    assert isinstance(result, Result)

tests/v1/unit/test_client.py

@@ -4,7 +4,6 @@
 from unittest.mock import patch
 
 import pytest
-import requests
 from aiohttp import ClientResponse
 from aiohttp import ClientResponseError
 from aiohttp import RequestInfo
@@ -392,30 +391,31 @@ async def test_post_to_field_endpoint_test_splits_multiple_fields() -> None:
         assert results == expected_data
 
 
-@patch("requests.post")
-def test_successful_post_to_sid_endpoint(
-    mock_post: Mock, test_client: PseudoClient
-) -> None:
-    mock_response = Mock(spec=requests.Response)
-    mock_response.status_code = 200
-    mock_response.raise_for_status.return_value = None
+@pytest.mark.asyncio
+async def test_successful_post_to_sid_endpoint(test_client: PseudoClient) -> None:
+    payload = [
+        {
+            "missing": ["magic", "sorcery"],
+            "datasetExtractionSnapshotTime": "2024-01-01T00:00:00Z",
+        }
+    ]
 
-    mock_post.return_value = mock_response
-    response = test_client._post_to_sid_endpoint(
-        path="test_path",
-        values=["value1", "value2"],
-    )
+    with patch("aiohttp.ClientSession.post", new_callable=AsyncMock) as mock_post:
+        mock_response = AsyncMock()
+        mock_response.__aenter__.return_value = mock_response
+        mock_response.json.return_value = payload
+        mock_post.return_value = mock_response
 
-    expected_json = {"fnrList": ["value1", "value2"]}
-    assert response == mock_response
-    mock_post.assert_called_once_with(
-        url="https://mocked.dapla-pseudo-service/test_path",
-        params=None,
-        headers={
-            "Authorization": "Bearer some-auth-token",
-            "X-Correlation-Id": ANY,
-        },
-        json=expected_json,
-        stream=True,
-        timeout=TIMEOUT_DEFAULT,
-    )
+        missing, snapshot_time = await test_client._post_to_sid_endpoint(
+            path="test_path",
+            values=["value1", "value2"],
+            sid_snapshot_date=None,
+        )
+
+        assert missing == ["magic", "sorcery"]
+        assert snapshot_time == "2024-01-01T00:00:00Z"
+
+        mock_post.assert_called_once()
+        _args, kwargs = mock_post.call_args
+        assert kwargs["json"] == {"fnrList": ["value1", "value2"]}
+        assert "X-Correlation-Id" in kwargs["headers"]

tests/v1/unit/test_validation.py

@@ -1,5 +1,5 @@
 from datetime import date
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock
 from unittest.mock import Mock
 from unittest.mock import patch
 
@@ -17,26 +17,20 @@
 
 
 @pytest_cases.fixture()
-def sid_lookup_missing_response() -> MagicMock:
-    mock_response = MagicMock()
-    mock_response.status_code = 200
-    mock_response.content = b'[{"missing": ["20859374701","01234567890"], "datasetExtractionSnapshotTime": "2023-08-31"}]'
-    return mock_response
+def sid_lookup_missing_response() -> tuple[list[str], str]:
+    return (["20859374701", "01234567890"], "2023-08-31")
 
 
 @pytest_cases.fixture()
-def sid_lookup_empty_response() -> MagicMock:
-    mock_response = MagicMock()
-    mock_response.status_code = 200
-    mock_response.content = b'[{"datasetExtractionSnapshotTime": "2023-08-31"}]'
-    return mock_response
+def sid_lookup_empty_response() -> tuple[list[str], str]:
+    return ([], "2023-08-31")
 
 
-@patch("dapla_pseudo.v1.PseudoClient._post_to_sid_endpoint")
+@patch("dapla_pseudo.v1.PseudoClient._post_to_sid_endpoint", new_callable=AsyncMock)
 def test_validate_with_full_response(
-    patched_post_to_sid_endpoint: Mock,
+    patched_post_to_sid_endpoint: AsyncMock,
     df_personer: pl.DataFrame,
-    sid_lookup_missing_response: MagicMock,
+    sid_lookup_missing_response: tuple[list[str], str],
 ) -> None:
     field_name = "fnr"
 
@@ -50,21 +44,20 @@ def test_validate_with_full_response(
     validation_df = validation_result.to_pandas()
     validation_metadata = validation_result.metadata_details
 
-    patched_post_to_sid_endpoint.assert_called_once_with(
-        "sid/lookup/batch",
-        ["11854898347", "01839899544", "16910599481"],
-        None,
-        stream=True,
+    patched_post_to_sid_endpoint.assert_awaited_once_with(
+        path="sid/lookup/batch",
+        values=["11854898347", "01839899544", "16910599481"],
+        sid_snapshot_date=None,
     )
     assert validation_df[field_name].tolist() == ["20859374701", "01234567890"]
     assert validation_metadata[field_name]["logs"] == ["SID snapshot time 2023-08-31"]
 
 
-@patch("dapla_pseudo.v1.PseudoClient._post_to_sid_endpoint")
+@patch("dapla_pseudo.v1.PseudoClient._post_to_sid_endpoint", new_callable=AsyncMock)
 def test_validate_with_empty_response(
-    patched_post_to_sid_endpoint: Mock,
+    patched_post_to_sid_endpoint: AsyncMock,
     df_personer: pl.DataFrame,
-    sid_lookup_empty_response: MagicMock,
+    sid_lookup_empty_response: tuple[list[str], str],
 ) -> None:
     field_name = "fnr"
 
@@ -79,10 +72,9 @@ def test_validate_with_empty_response(
     validation_metadata = validation_result.metadata_details
 
     patched_post_to_sid_endpoint.assert_called_once_with(
-        "sid/lookup/batch",
-        ["11854898347", "01839899544", "16910599481"],
-        date(2023, 8, 31),
-        stream=True,
+        path="sid/lookup/batch",
+        values=["11854898347", "01839899544", "16910599481"],
+        sid_snapshot_date=date(2023, 8, 31),
     )
     assert validation_df[field_name].tolist() == []
     assert validation_metadata[field_name]["logs"] == ["SID snapshot time 2023-08-31"]