Better handling of tokens near expiry

servicex/servicex_adapter.py

@@ -26,7 +26,7 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 import os
-from datetime import datetime
+import time
 from typing import Optional, Dict, List
 
 import httpx
@@ -69,8 +69,11 @@ def _get_bearer_token_file():
                 bearer_token = f.read().strip()
         return bearer_token
 
-    async def _get_authorization(self) -> Dict[str, str]:
-        if self.token:
+    async def _get_authorization(self, force_reauth: bool = False) -> Dict[str, str]:
+        now = time.time()
+        if (self.token and jwt.decode(self.token, verify=False)["exp"] - now > 60
+                and not force_reauth):
+            # if less than one minute validity, renew
             return {"Authorization": f"Bearer {self.token}"}
         else:
             bearer_token = self._get_bearer_token_file()
@@ -80,9 +83,8 @@ async def _get_authorization(self) -> Dict[str, str]:
             if not bearer_token and not self.refresh_token:
                 return {}
 
-            now = datetime.utcnow().timestamp()
-            if not self.token or \
-                    float(jwt.decode(self.token, verify=False)["exp"]) - now < 0:
+            if not self.token or force_reauth or\
+                    float(jwt.decode(self.token, verify=False)["exp"]) - now < 60:
                 await self._get_token()
             return {"Authorization": f"Bearer {self.token}"}
 
@@ -136,8 +138,7 @@ async def get_transform_status(self, request_id: str) -> TransformStatus:
         async with RetryClient(retry_options=retry_options) as client:
             try:
                 async for attempt in AsyncRetrying(
-                        retry=(retry_if_not_exception_type(AuthorizationError)
-                               | retry_if_not_exception_type(ValueError)),
+                        retry=retry_if_not_exception_type(ValueError),
                         stop=stop_after_attempt(3),
                         wait=wait_fixed(3),
                         reraise=True):
@@ -146,10 +147,18 @@ async def get_transform_status(self, request_id: str) -> TransformStatus:
                                               f"transformation/{request_id}",
                                               headers=headers) as r:
                             if r.status == 401:
+                                # perhaps we just ran out of auth validity the last time?
+                                # refetch auth then raise an error for retry
+                                headers = await self._get_authorization(True)
                                 raise AuthorizationError("Not authorized to access serviceX"
                                                          f"at {self.url}")
                             if r.status == 404:
                                 raise ValueError(f"Transform ID {request_id} not found")
+                            elif r.status > 400:
+                                o = await r.json()
+                                error_message = o.get('message', str(r))
+                                raise RuntimeError("ServiceX WebAPI Error during transformation: "
+                                                   f"{r.status} - {error_message}")
                             o = await r.json()
                             return TransformStatus(**o)
             except RuntimeError as e:

tests/test_servicex_adapter.py

@@ -33,6 +33,7 @@
 from pytest_asyncio import fixture
 from servicex.models import TransformRequest, ResultDestination, ResultFormat
 from servicex.servicex_adapter import ServiceXAdapter, AuthorizationError
+import time
 
 
 @fixture
@@ -186,7 +187,7 @@ async def test_get_transform_status(get, servicex, transform_status_response):
 
 @pytest.mark.asyncio
 @patch('servicex.servicex_adapter.RetryClient.get')
-async def test_get_transform_status_auth_error(get, servicex):
+async def test_get_transform_status_errors(get, servicex):
     with pytest.raises(AuthorizationError) as err:
         get.return_value.__aenter__.return_value.status = 401
         await servicex.get_transform_status("b8c508d0-ccf2-4deb-a1f7-65c839eebabf")
@@ -197,6 +198,15 @@ async def test_get_transform_status_auth_error(get, servicex):
         await servicex.get_transform_status("b8c508d0-ccf2-4deb-a1f7-65c839eebabf")
         assert "Transform ID b8c508d0-ccf2-4deb-a1f7-65c839eebabf not found" == str(err.value)
 
+    with pytest.raises(RuntimeError) as err:
+        get.return_value.__aenter__.return_value.status = 500
+
+        async def patch_json():
+            return {'message': 'fifteen'}
+        get.return_value.__aenter__.return_value.json = patch_json
+        await servicex.get_transform_status("b8c508d0-ccf2-4deb-a1f7-65c839eebabf")
+        assert "ServiceX WebAPI Error during transformation" in str(err.value)
+
 
 @pytest.mark.asyncio
 @patch('servicex.servicex_adapter.TransformStatus', side_effect=RuntimeError)
@@ -215,5 +225,13 @@ async def test_get_tranform_status_retry_error(get,
 @pytest.mark.asyncio
 async def test_get_authorization(servicex):
     servicex.token = "token"
-    r = await servicex._get_authorization()
-    assert r.get("Authorization") == "Bearer token"
+    servicex.refresh_token = "refresh"
+    with patch('google.auth.jwt.decode', return_value={'exp': time.time() + 90}):
+        r = await servicex._get_authorization()
+        assert r.get("Authorization") == "Bearer token"
+
+    with patch('servicex.servicex_adapter.ServiceXAdapter._get_token', return_value='token')\
+            as get_token:
+        with patch('google.auth.jwt.decode', return_value={'exp': time.time() - 90}):
+            r = await servicex._get_authorization()
+            get_token.assert_called_once()