Connect to Azure SQL database in SQL Alchemy using Entra ID tokens

We have a web application at work that needs to connect to our Azure SQL database for the duration of the request. To make querying the database easier, we use SQL Alchemy and pyodbc.

We have a few goals we want to achieve:

• Every request gets its own session. We want to open a new session when the request starts and close it once it is done.
• We want to connect to the database using the managed identity of the web service (Azure Function App or Azure Web App).
• We want to delegate the handling of the session and connection as much as we can.

Session lifetime management

First, to make sure we open a new session for every request, we can wrap the request functions into decorators that make sure a new session is created and subsequently destroyed.

To make sure our session is easily accessible from all over the application without having to hand it over in every function call, we use a singleton pattern. However, since we run multiple threads to handle multiple requests at the same time, we have to make sure that there are no race conditions with respect to the session objects.

SQL Alchemy has a great utility to make this easier: the scoped session.

To use this, we wrap our connection factory into a scoped_session() call:

from urllib.parse import quote_plus
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session

connection_string = "..."
engine = create_engine("mssql+pyodbc:///?odbc_connect={}".format(quote_plus(connection_string))
session_factory = sessionmaker(bind=engine)
Session = scoped_session(session_factory)

Now, whenever we want to use the session, we just call Session() and SQL Alchemy makes sure we reuse the existing session object. If we are done, we can call Session.remove() to close the session. There is no way for SQL Alchemy to tell if the thread is done so we have to do this ourselves.

Connect via Entra ID tokens using managed identity

There are some posts that describe how to setup the connection to Azure SQL databases using access tokens, but the best resource for this is the SQL Alchemy docs themselves.

Let's go over the details together. First, we need a connection string. Since we want to rely on managed identity (or the Azure CLI for local development), we do not put any credentials into the connection string:

Driver={ODBC Driver 18 for SQL Server};Database=YOUR_DB;Server=tcp:you.database.windows.net,1433;Encrypt=yes;TrustServerCertificate=no;Connection Timeout=30

Our general plan is:

• Add an event handler that fires whenever we connect to the database.
• Retrieve an access token whenever we connect to the database.
• Adjust the connection arguments so that we put the (fresh!) access token into the connection string.

Now, let's take a look at the code:

import struct
from urllib.parse import quote_plus

from azure.core.exceptions import ClientAuthenticationError
from azure.identity import DefaultAzureCredential
from sqlalchemy import event


@event.listens_for(engine, "do_connect")
def provide_token(dialect, conn_rec, cargs, cparams):
    # Uid may not be provided for the token authentication to work
    if "Uid=" in cargs[0]:
        return

    # Connection option for access tokens as defined in msodbcsql.h
    SQL_COPT_SS_ACCESS_TOKEN = 1256

    # remove the "Trusted_Connection" parameter that SQLAlchemy adds
    cargs[0] = cargs[0].replace(";Trusted_Connection=Yes", "")
    cargs[0] = cargs[0].replace(";Authentication=TokenIdentifiedPrincipal", "")

    # create token credential => try two ways to make local dev work
    scope = "https://database.windows.net/.default"
    try:
        cred = DefaultAzureCredential(exclude_managed_identity_credential=True).get_token(scope)
    except ClientAuthenticationError:
        cred = DefaultAzureCredential(
            exclude_managed_identity_credential=True,
            exclude_shared_token_cache_credential=True,
        ).get_token(scope)

    raw_token = cred.token.encode("utf-16-le")
    token_struct = struct.pack(f"<I{len(raw_token)}s", len(raw_token), raw_token)
    cparams["attrs_before"] = {SQL_COPT_SS_ACCESS_TOKEN: token_struct}

Summary

With this, we achieved our goals. With the scoped_session() we do not need to open a new session every time a request comes in (this will be handled for us) but we should close the session at the end so that we do not have too many dangling sessions.

We also connect to the SQL database using our own identity (for local dev) or the managed identity of the web service. We modify the connection string each time a new connection is created.

Acknowledgements

A big thank you to David for helping me figure out the concept of scoped sessions.