Search code examples
pythonapitokencontextmanager

Python rerun code with new token, when token has expired


I want to write some type of generic wrapper for api calls, that allows doing the requests without worrying about the token expiry, refreshes the token in the background.

Something like a context manager that handles token refresh in the background, invisible to the user. In order to do this, the "wrapper" must be able to re-run the code if a TokenExpiredException occured with the new Token.

For example, this code uses a 2 level try/except block, repeating the same call, and you have to pass api_call as string and duplicates code for calling the api:

def call_api_with_login(api_call: str, *args, **kwargs)
    """Call ``api_call`` method on the MyApi client with token error handling."""

    def get_method(client: ApiClient, call: str):
        """Get method from a dot-separated string"""
        return functools.reduce(getattr, call.split("."), client)

    api = MyApi()
    api_method = get_method(api.client, api_call)
    try:
        result = api_method(token, *args, **kwargs)
        report_api_call()
    except exceptions.TokenExpiredException as exc:
        token = api.login().token

        try:
            result = api_method(token, *args, **kwargs)
        except Exception as exc:
            logging.exception(exc)
            result = []

Besides the code duplication above and the fact that this "pattern" is quite limiting, it would be used like this:

call_api_with_login("books.list", author="Carl")

... which is kind of crappy, as we are passing method names as string, no access to code assistant, prone to errors, etc.

My initial idea is I would like to use something like a context manager to handle this, something like:

with authenticated_client as api_client, token:
    api_client.books.list(token, author="xyz")

The context manager would yield the client and token? ... However, there is no way I can think of to replay the inner code in case of an exception and refresh token (unless I do a loop of sorts in the context manager, more like a generator, maybe?)

def authenticated_client():

    api = MyApi()
    token = cache_session.cache_get("API_TOKEN")
    try:
        yield api, token
    except exceptions.TokenExpiredException as exc:
        token = api.login().token
        # ... how to rerun code?

Hope this example makes some sense without being fully descriptive of api client and all ...

Can someone recomend a better/cleaner way to do this, or maybe other ways to handle token refresh?

I tried the ideas explained above, the first works but is not really looking like good practice long term.


Solution

  • So what I mean by decorating is imagine you have

    def api_call(token, *args, **kwargs):
        ... # some logic here
    

    Your decorator will look something like this

    def authorize_on_expire(func):
        def wrapper(token, *args, **kwargs):
            try:
                result = func(token, *args, **kwargs)
            except exceptions.TokenExpiredException as e:
                token = ... # token refreshing logic
                result = func(token, *args, **kwargs)
            finally:
                return result
        return wrapper
    

    and you just decorate your api_call(...) like so:

    @authorize_on_expire
    def api_call(token, *args, **kwargs):
        ... # some logic here
    

    Context managers are created mostly for safely closing streams/connections/etc on error. One nice example that I have is rollback database transaction on any error and raise exception afterwards