如何在AWS托管的Airflow中使用联邦令牌连接BigQuery

问题 #

我需要在AWS托管的Airflow中构建读取Google BigQuery的连接器，DevOps团队提供了符合AIP-4117规范的外部账号凭证JSON：

{
  "type": "external_account",
  "audience": "xx",
  "subject_token_type": "xx",
  "service_account_impersonation_url": "xx",
  "token_uri": "xx",
  "credential_source": {"xx", "xx", "xx"}
}

该凭证已上传至AWS SSM参数存储，我能通过SsmHook.get_parameter_value()获取到它，但不知道如何配置BigQueryGetDataOperator完成认证。直接运行时报错：

airflow.exceptions.AirflowNotFoundException: The conn_id `google_cloud_default` isn't defined

解决方案 #

不需要创建Airflow的Google Cloud连接，直接通过SSM获取凭证并初始化BigQuery客户端，传递给Operator即可：

修改后的完整代码 #

from datetime import datetime
import json
from airflow import DAG
from airflow.providers.google.cloud.operators.bigquery import BigQueryGetDataOperator
from airflow.operators.empty import EmptyOperator
from airflow.providers.amazon.aws.hooks.ssm import SsmHook
from google.cloud import bigquery
from google.oauth2 import external_account

# 替换为实际项目、数据集、表名
PROJECT = "xxx"
DATASET = "xxx"
TABLE = "xxx"
# SSM参数路径
SSM_PARAM_NAME = "/prod/xxx/airflow"

def initialize_bigquery_client():
    # 从SSM拉取凭证并解析为字典
    ssm_hook = SsmHook()
    cred_json = ssm_hook.get_parameter_value(SSM_PARAM_NAME, with_decryption=True)
    cred_dict = json.loads(cred_json)
    
    # 基于AIP-4117凭证初始化认证信息
    credentials = external_account.Credentials.from_info(cred_dict)
    # 返回已认证的BigQuery客户端
    return bigquery.Client(project=PROJECT, credentials=credentials)

with DAG(
    "example-dag",
    schedule_interval=None,
    start_date=datetime(2021, 1, 1),
    catchup=False,
) as dag:
    start = EmptyOperator(task_id="start")
    
    get_data = BigQueryGetDataOperator(
        task_id="get_data",
        dataset_id=DATASET,
        table_id=TABLE,
        max_results=10,
        # 禁用默认的Google Cloud连接
        gcp_conn_id=None,
        # 传递自定义的已认证客户端
        hook_params={"client": initialize_bigquery_client()}
    )
    
    end = EmptyOperator(task_id="end")
    start >> get_data >> end

关键细节 #

• 绕过默认连接：设置gcp_conn_id=None，避免Operator尝试查找不存在的google_cloud_default连接。
• 凭证初始化：使用Google官方的external_account.Credentials.from_info()方法解析SSM获取的AIP-4117格式凭证，确保认证合规。
• 客户端传递：通过hook_params将已认证的BigQuery客户端直接传给Operator的底层Hook，确保操作使用正确的认证信息。

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