如何在Databricks FastAPI应用中传递Microsoft MFA至SQL Server托管实例

我用FastAPI构建了一个Databricks应用，用户在Azure Databricks完成Microsoft MFA认证后可访问该应用。应用需要连接支持Microsoft MFA的SQL Server托管实例（MI），我希望把已认证用户的凭据传递给应用，用于数据库连接。在开发环境中，我登录自己的Microsoft MFA账号后功能正常，但不确定在Databricks应用环境中能否实现这个需求。

代码片段 #

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"

    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):
        # 移除SQLAlchemy自动添加的"Trusted_Connection"参数
        cargs[0] = cargs[0].replace(";Trusted_Connection=Yes", "")

        # 生成令牌凭据
        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)

        # 将令牌应用到连接参数
        cparams["attrs_before"] = {SQL_COPT_SS_ACCESS_TOKEN: token_struct}

    # 测试连接
    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"通用错误：{str(e)}")

运行错误信息 #

通用错误：DefaultAzureCredential无法从包含的凭据中获取令牌。
尝试的凭据：
EnvironmentCredential：EnvironmentCredential认证不可用，环境变量未完全配置。
ManagedIdentityCredential：ManagedIdentityCredential认证不可用，IMDS端点无响应。
SharedTokenCacheCredential：SharedTokenCacheCredential认证不可用，缓存中未找到账户。
AzureCliCredential：未在路径中找到Azure CLI。
AzurePowerShellCredential：未安装PowerShell。
AzureDeveloperCliCredential：未找到Azure Developer CLI，请安装后使用'azd auth login'认证Azure账户。
如需缓解此问题，请查看DefaultAzureCredential排查指南。

解决方案 #

1. 利用Databricks用户上下文传递凭据 #

在Azure Databricks环境中，DefaultAzureCredential无法直接获取用户MFA凭据，需改用Databricks内置的用户身份传递机制：

• 通过databricks.sdk获取当前用户的会话令牌，再用该令牌请求Azure SQL Server的访问令牌，替代DefaultAzureCredential的令牌获取逻辑。
• 确保Databricks工作区已启用用户身份传递到Azure资源的功能，让用户的MFA身份能被下游服务识别。

2. 配置Databricks托管身份 #

若无需传递个人用户凭据，可使用Databricks工作区的托管身份：

• 为Databricks工作区启用系统分配/用户分配的托管身份。
• 给该托管身份授予SQL Server MI的db_datareader等必要权限。
• 修改代码，在DefaultAzureCredential中指定托管身份的客户端ID，确保ManagedIdentityCredential能正常获取令牌。

3. 验证环境依赖 #

确保Databricks运行环境中安装了ODBC Driver 18，同时SQL Server MI的防火墙规则允许Databricks工作区的出站IP访问。

内容的提问来源于stack exchange，提问作者Obirieni Simeo