Passing Microsoft MFA Auth from Databricks to MSSQL Managed Instance in a Databricks FastAPI App

[SObiero]

I have a Databricks App built using FastAPI. Users access this App after authenticating with Microsoft MFA on Databricks Azure Cloud. The App connects to an MSSQL Managed Instance (MI) that also supports Microsoft MFA.

I want the authenticated user's credentials to be passed to the App and used for database connections. In the development environment, this works fine since I am logged in with my Microsoft MFA. Can this be achieved in a Databricks App?

Below is a snippet of the code I am using with SQLAlchemy and PyODBC:

from sqlalchemy import create_engine, event, text
from azure.identity import DefaultAzureCredential
import struct
import pyodbc

try:
azure_credentials = DefaultAzureCredential()

TOKEN_URL = "https://database.windows.net/"
SQL_COPT_SS_ACCESS_TOKEN = 1256

server = 'azsql-mi-xyz.database.windows.net'
port = "1234"
database = 'database'
driver = "ODBC Driver 18 for SQL Server"

# Generate token
raw_token = azure_credentials.get_token(TOKEN_URL).token.encode("utf-16-le")

conn_str = f"mssql+pyodbc://@{server}:{port}/{database}?driver={driver}"
engine = create_engine(conn_str)

@event.listens_for(engine, "do_connect")
def provide_token(dialect, conn_rec, cargs, cparams):
# Remove "Trusted_Connection" parameter added by SQLAlchemy
cargs[0] = cargs[0].replace(";Trusted_Connection=Yes", "")

# Generate token credential
raw_token = azure_credentials.get_token(TOKEN_URL).token.encode("utf-16-le")
token_struct = struct.pack(f"<I{len(raw_token)}s", len(raw_token), raw_token)

# Apply it to connection parameters
cparams["attrs_before"] = {SQL_COPT_SS_ACCESS_TOKEN: token_struct}

# Test the connection
with engine.connect() as connection:
query = text("SELECT * FROM tablex")
rows = connection.execute(query).fetchall()

data = [dict(row._mapping) for row in rows]
return {"data": data}

except Exception as e:
raise HTTPException(status_code=500, detail=f"General Error: {str(e)}")