如何用SqlAlchemy通过令牌认证连接Azure Postgres实现AAD单点登录？

不用密码、基于AAD令牌的SqlAlchemy连接完全可以实现，核心是把AAD获取的令牌当作Postgres连接的"密码"传入，同时通过SqlAlchemy的扩展机制绕过连接字符串的限制。下面是两种实操方案：

方案一：利用SqlAlchemy的before_connect事件钩子 #

这个方法通过监听连接建立前的事件，动态注入AAD令牌作为密码，不需要修改连接字符串的结构。

步骤： #

1. 安装依赖：需要sqlalchemy、psycopg2-binary（或psycopg2）、azure-identity

pip install sqlalchemy psycopg2-binary azure-identity

2. 获取AAD访问令牌
   用Azure Identity库的DefaultAzureCredential（支持本地开发、Azure托管环境等多种场景的自动认证）获取针对Postgres的令牌：

from azure.identity import DefaultAzureCredential

def get_aad_token():
    credential = DefaultAzureCredential()
    # Postgres AAD认证的资源ID固定为https://ossrdbms-aad.database.windows.net
    token = credential.get_token("https://ossrdbms-aad.database.windows.net/.default")
    return token.token

3. 配置SqlAlchemy并绑定事件钩子

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.event import listens_for

# 基础连接字符串，不要包含password参数
conn_str = "postgresql+psycopg2://<AAD_USER_NAME>@<SERVER_NAME>.postgres.database.azure.com/<DATABASE_NAME>"

engine = create_engine(conn_str)

# 监听before_connect事件，注入令牌作为密码
@listens_for(Engine, "before_connect")
def inject_aad_token(dbapi_connection, connection_record):
    # 获取令牌
    token = get_aad_token()
    # 修改连接的密码属性
    dbapi_connection.password = token

# 测试连接
with engine.connect() as conn:
    result = conn.execute("SELECT version();")
    print(result.fetchone())

方案二：自定义连接创建函数 #

如果不想用事件钩子，可以直接自定义SqlAlchemy的连接创建逻辑，手动传入令牌作为密码：

from sqlalchemy import create_engine
from azure.identity import DefaultAzureCredential
import psycopg2

def get_aad_token():
    credential = DefaultAzureCredential()
    token = credential.get_token("https://ossrdbms-aad.database.windows.net/.default")
    return token.token

def custom_connect():
    token = get_aad_token()
    # 直接用psycopg2建立连接，传入令牌作为密码
    return psycopg2.connect(
        user="<AAD_USER_NAME>",
        host="<SERVER_NAME>.postgres.database.azure.com",
        dbname="<DATABASE_NAME>",
        password=token,
        sslmode="require"  # Azure Postgres强制SSL
    )

# 创建引擎时指定自定义连接函数
engine = create_engine("postgresql+psycopg2://", creator=custom_connect)

# 测试连接
with engine.connect() as conn:
    result = conn.execute("SELECT version();")
    print(result.fetchone())

关键注意事项 #

• AAD用户配置：Azure Postgres服务器必须启用AAD认证，且你使用的AAD身份（用户/服务主体/托管标识）已经被创建为Postgres中的用户，格式为CREATE USER "<AAD_PRINCIPAL_NAME>" WITH LOGIN;
• 令牌有效期：AAD令牌默认有效期约1小时，上述方案每次建立新连接时都会获取新令牌，自动处理过期问题
• SSL强制：Azure Postgres要求必须使用SSL连接，代码中确保sslmode="require"（psycopg2默认可能已经开启，但显式指定更稳妥）
• 驱动选择：推荐用psycopg2-binary避免编译依赖，如果你用psycopg2需要确保系统有Postgres开发库

内容的提问来源于stack exchange，提问作者Tien Dang