Merge pull request #469 from ssl-hep/fix-ClientPayloadError

Add retries on the `get` request of `get_transform_status` function

Author: kyungeonchoi

pyproject.toml

@@ -51,7 +51,8 @@ dependencies = [
     "make-it-sync",  # compatible versions controlled through func_adl
     "ruamel.yaml>=0.18",
     "ccorp-yaml-include-relative-path>=0.0.4",
-    "filelock>=3.12.0"
+    "filelock>=3.12.0",
+    "tenacity >= 9.0.0"
 ]
 
 [project.scripts]

servicex/servicex_adapter.py

@@ -32,6 +32,7 @@
 import httpx
 from aiohttp_retry import RetryClient, ExponentialRetry
 from google.auth import jwt
+from tenacity import AsyncRetrying, stop_after_attempt, wait_fixed, retry_if_not_exception_type
 
 from servicex.models import TransformRequest, TransformStatus
 
@@ -131,13 +132,26 @@ async def submit_transform(self, transform_request: TransformRequest):
 
     async def get_transform_status(self, request_id: str) -> TransformStatus:
         headers = await self._get_authorization()
-        retry_options = ExponentialRetry(attempts=5, start_timeout=10)
+        retry_options = ExponentialRetry(attempts=5, start_timeout=3)
         async with RetryClient(retry_options=retry_options) as client:
-            async with client.get(url=f"{self.url}/servicex/transformation/{request_id}",
-                                  headers=headers) as r:
-                if r.status == 401:
-                    raise AuthorizationError(f"Not authorized to access serviceX at {self.url}")
-                if r.status == 404:
-                    raise ValueError(f"Transform ID {request_id} not found")
-                o = await r.json()
-                return TransformStatus(**o)
+            try:
+                async for attempt in AsyncRetrying(
+                        retry=(retry_if_not_exception_type(AuthorizationError)
+                               | retry_if_not_exception_type(ValueError)),
+                        stop=stop_after_attempt(3),
+                        wait=wait_fixed(3),
+                        reraise=True):
+                    with attempt:
+                        async with client.get(url=f"{self.url}/servicex/"
+                                              f"transformation/{request_id}",
+                                              headers=headers) as r:
+                            if r.status == 401:
+                                raise AuthorizationError("Not authorized to access serviceX"
+                                                         f"at {self.url}")
+                            if r.status == 404:
+                                raise ValueError(f"Transform ID {request_id} not found")
+                            o = await r.json()
+                            return TransformStatus(**o)
+            except RuntimeError as e:
+                raise RuntimeError("ServiceX WebAPI Error "
+                                   f"while getting transform status: {e}")

tests/test_servicex_adapter.py

@@ -31,7 +31,6 @@
 import httpx
 import pytest
 from pytest_asyncio import fixture
-
 from servicex.models import TransformRequest, ResultDestination, ResultFormat
 from servicex.servicex_adapter import ServiceXAdapter, AuthorizationError
 
@@ -199,6 +198,20 @@ async def test_get_transform_status_auth_error(get, servicex):
         assert "Transform ID b8c508d0-ccf2-4deb-a1f7-65c839eebabf not found" == str(err.value)
 
 
+@pytest.mark.asyncio
+@patch('servicex.servicex_adapter.TransformStatus', side_effect=RuntimeError)
+@patch('servicex.servicex_adapter.RetryClient.get')
+async def test_get_tranform_status_retry_error(get,
+                                               mock_transform_status,
+                                               servicex,
+                                               transform_status_response):
+    with pytest.raises(RuntimeError) as err:
+        get.return_value.__aenter__.return_value.json.return_value = transform_status_response['requests'][0]  # NOQA: E501
+        get.return_value.__aenter__.return_value.status = 200
+        await servicex.get_transform_status("b8c508d0-ccf2-4deb-a1f7-65c839eebabf")
+        assert "ServiceX WebAPI Error while getting transform status:" in str(err.value)
+
+
 @pytest.mark.asyncio
 async def test_get_authorization(servicex):
     servicex.token = "token"