如何通过Operator自动获取GKE上Composer环境的内部信息？

推荐使用Google Cloud Provider专用工具 #

无需通过BashOperator调用gcloud，直接用Airflow自带的Google Cloud Provider包或GCP官方Python库，实现更优雅的信息获取，Composer环境默认已预装相关依赖。

方法1：基于Airflow Hooks实现 #

通过ComposerEnvironmentHook直接调用GCP API，无需外部命令：

from airflow.providers.google.cloud.hooks.composer import ComposerEnvironmentHook
from airflow.models.baseoperator import BaseOperator
from airflow.models import Variable

class FetchComposerInfoOperator(BaseOperator):
    def __init__(self, composer_env_name: str, composer_region: str, **kwargs):
        super().__init__(**kwargs)
        self.composer_env_name = composer_env_name
        self.composer_region = composer_region

    def execute(self, context):
        # 初始化Hook，默认使用google_cloud_default连接
        composer_hook = ComposerEnvironmentHook(gcp_conn_id="google_cloud_default")
        env_details = composer_hook.get_environment(
            environment_id=self.composer_env_name,
            region=self.composer_region
        )

        # 提取目标信息
        composer_info = {
            "COMPOSER_SERVICE_ACCOUNT": env_details["config"]["serviceAccount"],
            "COMPOSER_BUCKET": env_details["config"]["dagGcsPrefix"].split("/dags")[0].lstrip("gs://"),
            "COMPOSER_PROJECT": env_details["projectId"],
            "COMPOSER_PYTHON_VERSION": env_details["config"]["softwareConfig"]["pythonVersion"],
            "COMPOSER_VERSION": env_details["config"]["softwareConfig"]["imageVersion"].split("-composer-")[1],
            "COMPOSER_UI_URL": env_details["config"]["airflowUri"],
            "AIRFLOW_VERSION": env_details["config"]["softwareConfig"]["airflowVersion"]
        }

        # 存入Airflow变量和XCom
        for key, val in composer_info.items():
            Variable.set(key, val)
        self.xcom_push(context, key="composer_info", value=composer_info)
        return composer_info

使用示例

fetch_composer_info = FetchComposerInfoOperator(
    task_id="fetch_composer_details",
    composer_env_name="your-composer-environment-name",
    composer_region="your-region"
)

方法2：直接调用GCP Composer Python客户端 #

如果需要更灵活的控制，可使用google-cloud-composer库：

from google.cloud import composer_v1
from airflow.models.baseoperator import BaseOperator
from airflow.models import Variable

class FetchComposerInfoOperator(BaseOperator):
    def __init__(self, project_id: str, composer_env_name: str, composer_region: str, **kwargs):
        super().__init__(**kwargs)
        self.project_id = project_id
        self.composer_env_name = composer_env_name
        self.composer_region = composer_region

    def execute(self, context):
        client = composer_v1.EnvironmentsClient()
        env_path = client.environment_path(self.project_id, self.composer_region, self.composer_env_name)
        env_details = client.get_environment(name=env_path)

        composer_info = {
            "COMPOSER_SERVICE_ACCOUNT": env_details.config.service_account,
            "COMPOSER_BUCKET": env_details.config.dag_gcs_prefix.split("/dags")[0].lstrip("gs://"),
            "COMPOSER_PROJECT": self.project_id,
            "COMPOSER_PYTHON_VERSION": env_details.config.software_config.python_version,
            "COMPOSER_VERSION": env_details.config.software_config.image_version.split("-composer-")[1],
            "COMPOSER_UI_URL": env_details.config.airflow_uri,
            "AIRFLOW_VERSION": env_details.config.software_config.airflow_version
        }

        for key, val in composer_info.items():
            Variable.set(key, val)
        self.xcom_push(context, key="composer_info", value=composer_info)
        return composer_info

核心优势 #

• 避免解析gcloud命令输出的繁琐，稳定性更高
• 集成Airflow的变量和XCom系统，方便后续任务直接复用信息
• 自动继承Composer集群的服务账号权限，无需额外配置密钥

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