mdrs-client-python/mdrsclient/api/utils.py

import functools
from typing import Any, Callable, TypeVar, cast

from mdrsclient.api.users import UsersApi
from mdrsclient.connection import MDRSConnection
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:
    """
    Bring the access token up to date before a request goes out.

    Refreshing is a read-modify-write over a cache shared with every other client
    process on this machine, and a rotating provider stops honouring the refresh token
    it replaces. Two processes reaching this at once would otherwise both send the same
    token, and the loser would be left holding one the server no longer accepts, so the
    whole sequence runs under a lock that spans processes and the cache is re-read
    inside it.
    """
    with connection.lock:
        token = connection.token
        if token is None or not (token.is_refresh_required or token.is_expired):
            # Nothing to do, which is the answer for almost every request. The lock below
            # reaches across processes and is held for a round trip, so it is worth
            # knowing that before taking it.
            return
        with connection.cache_lock():
            connection.reload_cache()
            token = connection.token
            if token is None:
                return
            if token.is_expired:
                connection.logout()
                return
            if not token.is_refresh_required:
                return
            user_api = UsersApi(connection)
            try:
                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)