feat(auth): send a request again when its token lapsed in the queue

A request can wait to be served for longer than the access token it was
sent with lives, and comes back refused for a token that was valid when
it left. Uploads that take minutes make that wait ordinary.

- tell an expired token apart from any other refusal by the code the
  API returns, and retry only that one, once
- prefer a token another process left behind over minting a second: it
  saves a round trip, and a rotating provider would retire one that is
  still in use. Assign it rather than read it, since the setter is what
  rewrites the session header the retry will carry
- settle the token before the request so the one compared afterwards is
  the one that was actually sent
- report an overloaded server when the second attempt is refused too,
  and let API errors out of the upload path with their own type

mdrsclient/api/base.py

@@ -5,7 +5,13 @@ from pydantic import TypeAdapter
 from requests import Response
 
 from mdrsclient.connection import MDRSConnection
-from mdrsclient.exceptions import BadRequestException, ForbiddenException, UnauthorizedException, UnexpectedException
+from mdrsclient.exceptions import (
+    BadRequestException,
+    ForbiddenException,
+    TokenExpiredException,
+    UnauthorizedException,
+    UnexpectedException,
+)
 from mdrsclient.models.error import DRFStandardizedErrors
 
 
@@ -23,6 +29,11 @@ class BaseApi(ABC):
             if response.status_code == requests.codes.bad_request:
                 raise BadRequestException(errors.errors[0].detail)
             elif response.status_code == requests.codes.unauthorized:
+                # A request can wait in the server's queue for longer than the access
+                # token it was sent with lives, so an expired token here does not mean
+                # the session is over - it means this one request arrived too late.
+                if any(e.code == "token_not_valid" for e in errors.errors):
+                    raise TokenExpiredException("Access token expired before the request was served.")
                 raise UnauthorizedException("Login required.")
             elif response.status_code == requests.codes.forbidden:
                 raise ForbiddenException("You do not have enough permissions. Access is denied.")

mdrsclient/api/doi.py

@@ -4,7 +4,7 @@ from pydantic import TypeAdapter
 from pydantic.dataclasses import dataclass
 
 from mdrsclient.api.base import BaseApi
-from mdrsclient.api.utils import token_check
+from mdrsclient.api.utils import retry_on_expired_token, token_check
 from mdrsclient.models.doi import Doi
 
 
@@ -23,6 +23,7 @@ class DoiRetrieveResponse:
 class DoiApi(BaseApi):
     ENTRYPOINT: Final[str] = "v3/doi/"
 
+    @retry_on_expired_token
     def retrieve(self, doi_id: str) -> Doi:
         """Retrieve the folder associated with a DOI suffix ID (GET v3/doi/{id}/)."""
         url = self.ENTRYPOINT + doi_id + "/"

mdrsclient/api/files.py

@@ -8,8 +8,8 @@ from pydantic.dataclasses import dataclass
 from requests_toolbelt.multipart.encoder import MultipartEncoder
 
 from mdrsclient.api.base import BaseApi
-from mdrsclient.api.utils import token_check
+from mdrsclient.api.utils import retry_on_expired_token, token_check
-from mdrsclient.exceptions import UnexpectedException
+from mdrsclient.exceptions import MDRSException, UnexpectedException
 from mdrsclient.models import File
 
 
@@ -30,6 +30,7 @@ class FilesApi(BaseApi):
     ENTRYPOINT: Final[str] = "v3/files/"
     FALLBACK_MIMETYPE: Final[str] = "application/octet-stream"
 
+    @retry_on_expired_token
     def list(self, folder_id: str, page_num: int) -> FilesApiListResponse:
         url = self.ENTRYPOINT
         token_check(self.connection)
@@ -38,6 +39,7 @@ class FilesApi(BaseApi):
         self._raise_response_error(response)
         return TypeAdapter(FilesApiListResponse).validate_python(response.json())
 
+    @retry_on_expired_token
     def retrieve(self, id: str) -> File:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + id + "/"
@@ -46,6 +48,7 @@ class FilesApi(BaseApi):
         self._raise_response_error(response)
         return TypeAdapter(File).validate_python(response.json())
 
+    @retry_on_expired_token
     def create(self, folder_id: str, path: str) -> str:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT
@@ -66,10 +69,14 @@ class FilesApi(BaseApi):
             raise UnexpectedException(f"Could not open `{path}` file.")
         except MemoryError:
             raise UnexpectedException("Out of memory.")
+        except MDRSException:
+            # Already says what went wrong, and the caller may want to act on the kind.
+            raise
         except Exception as e:
             raise UnexpectedException("Unspecified error.") from e
         return ret.id
 
+    @retry_on_expired_token
     def update(self, file: File, path: str | None) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + file.id + "/"
@@ -96,6 +103,7 @@ class FilesApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def destroy(self, file: File) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + file.id + "/"
@@ -104,6 +112,7 @@ class FilesApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def move(self, file: File, folder_id: str, name: str) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + file.id + "/move/"
@@ -113,6 +122,7 @@ class FilesApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def copy(self, file: File, folder_id: str, name: str) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + file.id + "/copy/"
@@ -122,6 +132,7 @@ class FilesApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def metadata(self, file: File) -> dict[str, Any]:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + file.id + "/metadata/"
@@ -130,6 +141,7 @@ class FilesApi(BaseApi):
         self._raise_response_error(response)
         return response.json()
 
+    @retry_on_expired_token
     def download(self, file: File, path: str) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = file.download_url

mdrsclient/api/folders.py

@@ -5,7 +5,7 @@ from pydantic import TypeAdapter
 from pydantic.dataclasses import dataclass
 
 from mdrsclient.api.base import BaseApi
-from mdrsclient.api.utils import token_check
+from mdrsclient.api.utils import retry_on_expired_token, token_check
 from mdrsclient.exceptions import UnauthorizedException
 from mdrsclient.models import Folder, FolderSimple
 
@@ -18,6 +18,7 @@ class FoldersApiCreateResponse:
 class FoldersApi(BaseApi):
     ENTRYPOINT: Final[str] = "v3/folders/"
 
+    @retry_on_expired_token
     def list(self, laboratory_id: int, path: str) -> list[FolderSimple]:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT
@@ -30,6 +31,7 @@ class FoldersApi(BaseApi):
             ret.append(TypeAdapter(FolderSimple).validate_python(data))
         return ret
 
+    @retry_on_expired_token
     def retrieve(self, id: str) -> Folder:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + id + "/"
@@ -39,6 +41,7 @@ class FoldersApi(BaseApi):
         ret = TypeAdapter(Folder).validate_python(response.json())
         return ret
 
+    @retry_on_expired_token
     def create(self, name: str, parent_id: str) -> str:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT
@@ -49,6 +52,7 @@ class FoldersApi(BaseApi):
         ret = TypeAdapter(FoldersApiCreateResponse).validate_python(response.json())
         return ret.id
 
+    @retry_on_expired_token
     def update(self, folder: FolderSimple) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + folder.id + "/"
@@ -61,6 +65,7 @@ class FoldersApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def destroy(self, id: str, recursive: bool) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + id + "/"
@@ -70,6 +75,7 @@ class FoldersApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def auth(self, id: str, password: str) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + id + "/auth/"
@@ -81,6 +87,7 @@ class FoldersApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def acl(self, id: str, access_level: int, recursive: bool, password: str | None) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + id + "/acl/"
@@ -94,6 +101,7 @@ class FoldersApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def move(self, folder: FolderSimple, folder_id: str, name: str) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + folder.id + "/move/"
@@ -103,6 +111,7 @@ class FoldersApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def copy(self, folder: FolderSimple, folder_id: str, name: str) -> bool:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + folder.id + "/copy/"
@@ -112,6 +121,7 @@ class FoldersApi(BaseApi):
         self._raise_response_error(response)
         return True
 
+    @retry_on_expired_token
     def metadata(self, id: str) -> dict[str, Any]:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT + id + "/metadata/"

mdrsclient/api/laboratories.py

@@ -3,13 +3,14 @@ from typing import Final
 from pydantic import TypeAdapter
 
 from mdrsclient.api.base import BaseApi
-from mdrsclient.api.utils import token_check
+from mdrsclient.api.utils import retry_on_expired_token, token_check
 from mdrsclient.models import Laboratories, Laboratory
 
 
 class LaboratoriesApi(BaseApi):
     ENTRYPOINT: Final[str] = "v3/laboratories/"
 
+    @retry_on_expired_token
     def list(self) -> Laboratories:
         # print(self.__class__.__name__ + "::" + sys._getframe().f_code.co_name)
         url = self.ENTRYPOINT

mdrsclient/api/users.py

@@ -8,7 +8,8 @@ from mdrsclient.api.base import BaseApi
 from mdrsclient.exceptions import UnauthorizedException
 from mdrsclient.models import Token, User
 
-# (connect, read) seconds for the token refresh.
+# (connect, read) seconds for the token refresh. Uploads are served by a separate
+# instance, so this one is not queued behind them and has no reason to be slow.
 TOKEN_REFRESH_TIMEOUT: Final[tuple[float, float]] = (5.0, 30.0)
 
 

mdrsclient/api/utils.py

@@ -1,6 +1,22 @@
+import functools
+from typing import Any, Callable, TypeVar, cast
+
 from mdrsclient.api.users import UsersApi
 from mdrsclient.connection import MDRSConnection
-from mdrsclient.exceptions import UnauthorizedException
+from mdrsclient.exceptions import (
+    MDRSException,
+    ServerBusyException,
+    TokenExpiredException,
+    UnauthorizedException,
+)
+from mdrsclient.models import Token
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+SERVER_BUSY_MESSAGE = (
+    "The server took too long to start handling the request and may be overloaded. "
+    "Try again, or reduce the number of parallel transfers."
+)
 
 
 def token_check(connection: MDRSConnection) -> None:
@@ -36,3 +52,73 @@ def token_check(connection: MDRSConnection) -> None:
                 connection.token = user_api.tokenRefresh(token)
             except UnauthorizedException:
                 connection.logout()
+
+
+def token_recover(connection: MDRSConnection, used: Token | None) -> bool:
+    """
+    Get a usable access token after one was refused for having expired.
+
+    A request can sit in the server's queue for longer than its access token lives, and
+    while it waited another thread or process may already have refreshed. Prefer what
+    they left behind: refreshing again would spend a round trip, and with a provider
+    that rotates refresh tokens it would retire one that is still in use.
+
+    Returns whether the caller now holds a token worth retrying with.
+    """
+    with connection.lock, connection.cache_lock():
+        connection.reload_cache()
+        token = connection.token
+        if token is None:
+            return False
+        if used is None or token != used:
+            # Somebody else has been here since the request went out. Assign rather than
+            # just read: the setter is what rewrites the session's Authorization header,
+            # so a plain reload would leave the retry carrying the refused token.
+            connection.token = token
+            return True
+        user_api = UsersApi(connection)
+        try:
+            connection.token = user_api.tokenRefresh(token)
+        except UnauthorizedException:
+            connection.logout()
+            return False
+        return True
+
+
+def retry_on_expired_token(func: F) -> F:
+    """
+    Send a request again once when the server found its access token expired.
+
+    The token is checked before every request, so this only happens when the request was
+    held long enough on the way in for a valid token to lapse - a queue behind uploads
+    that take minutes, most often.
+    """
+
+    @functools.wraps(func)
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        connection: MDRSConnection = self.connection
+        # Settle the token first so `used` is what the request actually carries. The
+        # wrapped method checks it again, which costs nothing once it is already current.
+        token_check(connection)
+        used = connection.token
+        try:
+            return func(self, *args, **kwargs)
+        except TokenExpiredException:
+            try:
+                recovered = token_recover(connection, used)
+            except MDRSException:
+                raise
+            except Exception as e:
+                # A busy server queues the refresh too, and failing here would report an
+                # endpoint the user never asked for rather than the reason.
+                raise ServerBusyException(SERVER_BUSY_MESSAGE) from e
+            if not recovered:
+                raise
+        try:
+            return func(self, *args, **kwargs)
+        except TokenExpiredException as e:
+            # A second refusal, with a token that was current when it left. The session is
+            # fine; the server is not starting requests before their credentials lapse.
+            raise ServerBusyException(SERVER_BUSY_MESSAGE) from e
+
+    return cast(F, wrapper)

mdrsclient/exceptions.py

@@ -28,6 +28,18 @@ class UnauthorizedException(MDRSException):
     pass
 
 
+class TokenExpiredException(UnauthorizedException):
+    """Thrown when the access token was still valid when sent but had expired on arrival"""
+
+    pass
+
+
+class ServerBusyException(MDRSException):
+    """Thrown when the server did not start handling a request before its token lapsed"""
+
+    pass
+
+
 class ForbiddenException(MDRSException):
     """Thrown when the current user does not have enough privileges to access the resource"""
 

tests/test_token_cache.py

@@ -5,11 +5,11 @@ from unittest.mock import patch
 
 import jwt
 
-from mdrsclient.api.utils import token_check
+from mdrsclient.api.utils import retry_on_expired_token, token_check, token_recover
 from mdrsclient.cache import CacheFile
 from mdrsclient.config import normalize_url
 from mdrsclient.connection import MDRSConnection
-from mdrsclient.exceptions import IllegalArgumentException
+from mdrsclient.exceptions import IllegalArgumentException, ServerBusyException, TokenExpiredException
 from mdrsclient.models import Token
 
 REMOTE = "unittest"
@@ -112,5 +112,87 @@ class TestUrlNormalization(unittest.TestCase):
                 normalize_url(url)
 
 
+class TestExpiredTokenRetry(unittest.TestCase):
+    """A request can wait in the server's queue for longer than its access token lives."""
+
+    def setUp(self) -> None:
+        self.tmp = tempfile.TemporaryDirectory()
+        patcher = patch("mdrsclient.cache.CONFIG_DIRNAME", self.tmp.name)
+        patcher.start()
+        self.addCleanup(patcher.stop)
+        self.addCleanup(self.tmp.cleanup)
+        self.connection = MDRSConnection(REMOTE, "http://localhost:8000/api")
+        self.connection.token = make_token(3600, 86400, "sent")
+
+    def test_a_token_someone_else_refreshed_is_reused_rather_than_replaced(self) -> None:
+        used = self.connection.token
+        rotated = make_token(3600, 86400, "rotated")
+        CacheFile(REMOTE).token = rotated  # another process, while the request waited
+
+        with patch("mdrsclient.api.utils.UsersApi") as users_api:
+            recovered = token_recover(self.connection, used)
+
+        self.assertTrue(recovered)
+        users_api.assert_not_called()
+        self.assertEqual(self.connection.token, rotated)
+        # The retry carries whatever the session header says, not what the cache holds,
+        # so reading the new token through is not enough on its own.
+        self.assertEqual(self.connection.session.headers["Authorization"], f"Bearer {rotated.access}")
+
+    def test_the_token_is_refreshed_when_nobody_else_has(self) -> None:
+        used = self.connection.token
+        rotated = make_token(3600, 86400, "rotated")
+
+        with patch("mdrsclient.api.utils.UsersApi") as users_api:
+            users_api.return_value.tokenRefresh.return_value = rotated
+            recovered = token_recover(self.connection, used)
+
+        self.assertTrue(recovered)
+        users_api.return_value.tokenRefresh.assert_called_once()
+        self.assertEqual(self.connection.token, rotated)
+
+    def test_the_request_is_sent_once_more_and_only_once(self) -> None:
+        calls: list[str] = []
+
+        class Api:
+            def __init__(self, connection: MDRSConnection) -> None:
+                self.connection = connection
+
+            @retry_on_expired_token
+            def send(self) -> str:
+                calls.append("sent")
+                if len(calls) == 1:
+                    raise TokenExpiredException("expired on arrival")
+                return "ok"
+
+        with patch("mdrsclient.api.utils.UsersApi") as users_api:
+            users_api.return_value.tokenRefresh.return_value = make_token(3600, 86400, "rotated")
+            self.assertEqual(Api(self.connection).send(), "ok")
+
+        self.assertEqual(len(calls), 2)
+
+    def test_a_second_refusal_is_reported_as_an_overloaded_server(self) -> None:
+        calls: list[str] = []
+
+        class Api:
+            def __init__(self, connection: MDRSConnection) -> None:
+                self.connection = connection
+
+            @retry_on_expired_token
+            def send(self) -> str:
+                calls.append("sent")
+                raise TokenExpiredException("expired on arrival")
+
+        with patch("mdrsclient.api.utils.UsersApi") as users_api:
+            users_api.return_value.tokenRefresh.return_value = make_token(3600, 86400, "rotated")
+            with self.assertRaises(ServerBusyException) as caught:
+                Api(self.connection).send()
+
+        # Sent twice and no more, and the message names the cause the user can act on
+        # rather than an authentication failure that would send them to the login form.
+        self.assertEqual(len(calls), 2)
+        self.assertIn("overloaded", str(caught.exception))
+
+
 if __name__ == "__main__":
     unittest.main()