Implementing Access Token Refreshes in Python

Learn how to implement access token refreshes for Airbyte API applications using Python

DATA

November 27, 2024

10 min read

Quinton Wall

Summarize with AI:

Implementing Access Token Refreshes in Python

I’ve been working on a sample Python app that demonstrates how to use the Airbyte APIs to perform typical activities such as connecting to workspaces, getting a list of connections, and checking whether they are active….basically common functions you may need for CI/CD testing of data pipelines.

When working with Airbyte’s Python libraries, things can be a little confusing at first glance. Basically, you have two options:

API Client libraries
1. The API client libraries, available via pip for Python development, provide complete coverage of the APIs. Documentation for APIs is great and always you to quickly test endpoints, and generate code snippets for cURL, Ruby, Node, and Python. Beyond the API docs, Python examples can be found on the Pypi page. These tests and snippets work OOTB for Cloud, but for OSS, the URI is unique. Eg:
  1. Cloud: https://api.airbyte.com/v1/sources
  2. OSS: https://your-instance:port/api/v1/sources
PyAirbyte
1. PyAirbyte is an open-source library that packages Airbyte connectors and makes them available in Python. It is compatible with Airbyte Cloud and Open Source. PyAirbyte is still relatively new, with the vision to provide full API coverage, but it’s not fully there yet. What is does cover, and how it is used, is vastly improved over the existing API client libraries. If you are comfortable using Pandas, and typical Python coding practices, PyAirbyte is for you. Documentation is here and targeted more for Python devs via Collab notebook samples vs. inline calling functionality like the API docs. Wherever possible, I would recommend always starting with PyAirbtye. It’s the future of our Python libraries.

Regardless of which route you choose you will need to manage access token expiration in your app. Airbyte tokens expire after 180 seconds. Due to the short duration, it is best practice to always refresh your token before an API call. There are a number of ways you can achieve this. The airbyte-api has a schemaclientcredentials model with an example on how to use it here, but to keep things really simple, I’m going to hit the API endpoint directly. In my case, I like to use either the reference site or Postman to verify endpoints are working before adding them to my app.

I’m building my app using Streamlit (import streamlit as st), using python3.12. I wrote a function which uses session state to manage tokens for some basic token caching, and to perform a token refresh every 120 seconds:


def refresh_airbyte_token():
    try:
         # Check if we recently refreshed
        last_refresh = st.session_state.get('last_token_refresh', 0)
        if time.time() - last_refresh < 120:  # 120 seconds = 2 minutes
            return st.session_state.get('current_token')
        
        #if older, go get a new one
        st.info("Fetching an updated token.")
        refresh_token = os.getenv('AIRBYTE_REFRESH_TOKEN')
        if not refresh_token:
            raise ValueError("AIRBYTE_REFRESH_TOKEN not found in environment variables")

        base_url = os.getenv('AIRBYTE_HOST').rstrip('/')
        response = requests.post(
            f'{base_url}/v1/applications/token',
            data={
                'grant_type': 'refresh_token',
                'client_id': os.getenv('AIRBYTE_CLIENT_ID'),
                'client_secret': os.getenv('AIRBYTE_CLIENT_SECRET'),
                'refresh_token': refresh_token
            }
        )
        response.raise_for_status()
        
        new_token = response.json().get('access_token')
        if not new_token:
            raise ValueError("No access token in response")

        # Update session state with new clients
        st.session_state.workspace = cloud.CloudWorkspace(
            os.getenv('AIRBYTE_WORKSPACE_ID'),
            new_token
        )
        st.session_state.client = AirbyteAPI(
            security=models.Security(bearer_auth=new_token)
        )
        st.session_state.last_token_refresh = time.time()
        st.session_state.current_token = new_token
        
        return new_token

    except Exception as e:
        st.error(f"Failed to refresh token: {str(e)}")
        raise

There’s a lot going on here. Let’s break it down:

‍

Check for last refresh & obtain keys


 last_refresh = st.session_state.get('last_token_refresh', 0)
        if time.time() - last_refresh < 120:  # 120 seconds = 2 minutes
            return st.session_state.get('current_token')

I’m using a variable, last_refresh, that is stored in session state. If it’s been more than 120 seconds, it’s time to refresh the token. Because I don’t want to hard code values such as client_id and client_secret, I can leverage a library such as python-dotenv to store values in a .env file.

If you need to obtain your Airbyte client_id, client_secret, and access token, you obtain these via Settings > Applications in the Airbyte UI.

‍

Fetch your application credentials via the Airbyte UI

Request a token


refresh_token = os.getenv('AIRBYTE_REFRESH_TOKEN')
        if not refresh_token:
            raise ValueError("AIRBYTE_REFRESH_TOKEN not found in environment variables")

        base_url = os.getenv('AIRBYTE_HOST').rstrip('/')
        response = requests.post(
            f'{base_url}/v1/applications/token',
            data={
                'grant_type': 'refresh_token',
                'client_id': os.getenv('AIRBYTE_CLIENT_ID'),
                'client_secret': os.getenv('AIRBYTE_CLIENT_SECRET'),
                'refresh_token': refresh_token
            }
        )
        response.raise_for_status()
        
        new_token = response.json().get('access_token')
        if not new_token:
            raise ValueError("No access token in response")

You have to call the Airbyte API directly, hitting https://api.airbyte.io/v1/application/token with a grant_type of refresh_token. Check out the docs here if you want to test it out before adding to your code. Keep in mind, this URL is different for cloud vs. OSS. In my example, I am using Cloud.

Update session state with new clients


 # Update session state with new clients
        st.session_state.workspace = cloud.CloudWorkspace(
            os.getenv('AIRBYTE_WORKSPACE_ID'),
            new_token
        )
        st.session_state.client = AirbyteAPI(
            security=models.Security(bearer_auth=new_token)
        )
        st.session_state.last_token_refresh = time.time()
        st.session_state.current_token = new_token
        
        return new_token

our code might look a little different here. I’m actually using both the API client wrapper and PyAirbyte in my complete app. Why? PyAirbyte doesn’t have full coverage yet for some features I needed. TL:DR; I need to update two clients and store them in session. You may have to just store one.

Now, with my refresh_airbyte_token() function set up, whenever I call the Airbyte APIs, I can make sure I have a valid token:


refresh_airbyte_token()
connection = st.session_state.workspace.get_connection(row['Connection ID'])

That’s about it. Adding this refresh logic to your app or python script whenever you are interacting with the Airbyte APIs is critical to ensuring your processes performs correctly. I’m still putting the finishing touches to the full app. Make sure you check out the project repo in GitHub to star and watch and keep up to date.

About the Author

Quinton Wall

Quinton is the Head of Developer Relations at Airbyte