SageMaker JupyterLab中Snowflake Entra ID OAuth客户端凭证流连接实践

前置准备（Snowflake侧） #

• 在Snowflake中创建OAuth安全集成：
  执行CREATE SECURITY INTEGRATION语句，指定TYPE = OAUTH，配置OAUTH_CLIENT_CREDENTIALS相关参数，关联目标权限角色，并设置允许的微软AD令牌端点。完成后记录生成的OAUTH_CLIENT_ID和OAUTH_CLIENT_SECRET（务必保密）。
• 确保集成关联的角色拥有访问指定仓库、数据库、Schema的权限。

SageMaker JupyterLab实现步骤 #

1. 安装依赖库 #

!pip install snowflake-connector-python requests

2. 获取微软AD OAuth访问令牌 #

编写代码请求微软令牌端点，获取可用于Snowflake的访问令牌：

import requests
import os

# 从环境变量或SageMaker Secrets Manager读取敏感信息（禁止硬编码）
client_id = os.getenv("SNOWFLAKE_OAUTH_CLIENT_ID")
client_secret = os.getenv("SNOWFLAKE_OAUTH_CLIENT_SECRET")
tenant_id = os.getenv("AZURE_TENANT_ID")
snowflake_account = "<account_identifier>"  # 示例：xy12345

# 微软AD令牌请求端点
token_url = f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"

# 请求参数：scope必须匹配Snowflake账户域名
payload = {
    "grant_type": "client_credentials",
    "client_id": client_id,
    "client_secret": client_secret,
    "scope": f"https://{snowflake_account}.snowflakecomputing.com/.default"
}

# 发送请求并处理响应
response = requests.post(token_url, data=payload)
response.raise_for_status()  # 抛出请求错误便于调试
access_token = response.json()["access_token"]

3. 用令牌连接Snowflake #

import snowflake.connector

# 建立连接
ctx = snowflake.connector.connect(
    account=snowflake_account,
    host=f"{snowflake_account}.snowflakecomputing.com",
    authenticator="OAUTH_CLIENT_CREDENTIALS",
    token=access_token,
    warehouse="test_warehouse",
    database="test_db",
    schema="test_schema"
)

# 测试连接有效性
cursor = ctx.cursor()
cursor.execute("SELECT CURRENT_VERSION()")
print(f"Snowflake版本：{cursor.fetchone()[0]}")

# 关闭资源
cursor.close()
ctx.close()

关键参数说明 #

• authenticator="OAUTH_CLIENT_CREDENTIALS"：强制启用客户端凭证OAuth流，是触发该认证方式的核心参数。
• token：传入从微软AD获取的访问令牌，无需在Snowflake连接参数中直接传入oauth_client_id或oauth_client_secret（这两个参数仅用于向微软请求令牌，而非Snowflake连接参数）。
• account/host：必须与请求令牌时scope中的Snowflake账户域名完全一致。
• user参数：客户端凭证流中无需指定（权限由Snowflake OAuth集成关联的角色决定），若指定需确保该用户与集成角色匹配。

SageMaker安全最佳实践 #

避免硬编码敏感信息，使用SageMaker Secrets Manager存储并读取：

import boto3
import json

secrets_manager = boto3.client("secretsmanager")
secret = secrets_manager.get_secret_value(SecretId="snowflake-oauth-secrets")
secret_dict = json.loads(secret["SecretString"])

client_id = secret_dict["client_id"]
client_secret = secret_dict["client_secret"]

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