How To Refresh an Access Token Using Decorators

How To Refresh an Access Token Using Decorators

When I was creating a one script, which uses JWT for authentication, I began to run into challenges with longer scripts that extended beyond the expiration of a single JWT.

To elegantly solve this, I used decorators to check the token’s expiration and request a new token if necessary. This article goes over the framework I set up so that you can apply a similar mechanism in your own scripts.

Setting the Scene


To get started, I’ve outlined all the parts needed to set up our token-refreshing framework.

import time
import requests

class myAPI():
    host = None
    key = None
    secret = None
    access_token = None
    access_token_expiration = None

    def __init__(self,host,key,secret):
        # the function that is executed when
        # an instance of the class is created
        pass

    def getAccessToken(self):
        # the function that is 
        # used to request the JWT
        pass

    class Decorators():
        @staticmethod
        def refreshToken(decorated):
            # the function that is used to check
            # the JWT and refresh if necessary
            pass

Our class will work by requiring the fields necessary to request the JWT. Without access, there is nothing more our class will do.

Writing the Access Token Function


Fleshing out getAccessToken() is going to be dependent on whatever API you are trying to interact with. For that reason, I won’t include any code constructing and executing the JWT request.

def getAccessToken(self):
    # the function that is 
    # used to request the JWT
    try:
        # build the JWT and store
        # in the variable `token_body`

        # request an access token
        request = requests.post(self.host,data=token_body)

        # optional: raise exception for status code
        request.raise_for_status()
    except Exception as e:
        print(e)
        return None
    else:
        # assuming the response's structure is
        # {"access_token": ""}
        return request.json()['access_token']

I highly recommend using a try/except/else along with raise_for_status() for requesting the access token. As we’re working with APIs, the request can fail and not return an access token but it can be successful in the eyes of Python.

Writing the Initialization Function


The __init__() function will be executed whenever we create a new instance of our myAPI class. What we want to do in this function is request an access token and, if successful, set the expiration time.

def __init__(self,host,key,secret):
    # the function that is executed when
    # an instance of the class is created
    self.host = host
    self.key = key
    self.secret = secret

    try:
        self.access_token = self.getAccessToken()
        if self.access_token is None:
            raise Exception("Request for access token failed.")
    except Exception as e:
        print(e)
    else:
        self.access_token_expiration = time.time() + 3500

We make the assumption that the access token should be returned and stored in self.access_token, so, if the value is None after the request, we raise an exception.

Otherwise, we set the expiration time for our access token. I like to give the class a small buffer, so if my token expires in one hour (3,600 seconds) I’m going to set the expiration for 3,500 seconds.

Writing the Decorator Class and Function


Decorators are an easy way of wrapping a function within another.

If you’re not familiar with decorators, I would recommend checking out the primer on Real Python. For our API class, we want to wrap all API functions with an access token check-and-refresh.

class Decorators():
    @staticmethod
    def refreshToken(decorated):
        # the function that is used to check
        # the JWT and refresh if necessary
        def wrapper(api,*args,**kwargs):
            if time.time() > api.access_token_expiration:
                api.getAccessToken()
            return decorated(api,*args,**kwargs)

        return wrapper

Putting It All Together and Using the Decorator


Now that we’re all set, it’s time to create a new function that will use the decorator. All you need to do is put @Decorators.refreshToken above any functions that will make an API request.

I’ve put the code all together along with an empty sample function at the end.

class myAPI():
    host = None
    key = None
    secret = None
    access_token = None
    access_token_expiration = None

    def __init__(self,host,key,secret):
        # the function that is executed when
        # an instance of the class is created
        self.host = host
        self.key = key
        self.secret = secret

        try:
            self.access_token = self.getAccessToken()
            if self.access_token is None:
                raise Exception("Request for access token failed.")
        except Exception as e:
            print(e)
        else:
            self.access_token_expiration = time.time() + 3500

    def getAccessToken(self):
        # the function that is 
        # used to request the JWT
        try:
            # build the JWT and store
            # in the variable `token_body`

            # request an access token
            request = requests.post(self.host,data=token_body)

            # optional: raise exception for status code
            request.raise_for_status()
        except Exception as e:
            print(e)
            return None
        else:
            # assuming the response's structure is
            # {"access_token": ""}
            return request.json()['access_token']

    class Decorators():
        @staticmethod
        def refreshToken(decorated):
            # the function that is used to check
            # the JWT and refresh if necessary
            def wrapper(api,*args,**kwargs):
                if time.time() > api.access_token_expiration:
                    api.getAccessToken()
                return decorated(api,*args,**kwargs)

            return wrapper

    @Decorators.refreshToken
    def someRequest():
        # make our API request
        pass

I hope you enjoyed this tutorial and it serves you well. This intelligent refresh mechanism is far superior to arbitrarily requesting new JWTs mid-script.