Dataflow工作器中环境变量配置的最佳实践

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions, StandardOptions
import argparse
import os

class CustomPipelineOptions(PipelineOptions):
    """
    自定义管道选项，用于传递应用程序特定参数。
    """
    @classmethod
    def _add_argparse_args(cls, parser):
        super()._add_argparse_args(parser)
        parser.add_argument(
            '--otel_service_name',
            dest='otel_service_name',
            default='default-service',
            help='OpenTelemetry服务名称。'
        )
        parser.add_argument(
            '--otel_resource_attributes',
            dest='otel_resource_attributes',
            default='key1=value1,key2=value2',
            help='OpenTelemetry资源属性，格式为key=value,key2=value2。'
        )
        # 您可以根据需要添加更多自定义参数

class ProcessBillRequests:
    class FetchBillInformation(beam.DoFn):
        def setup(self):
            # 在DoFn初始化时获取管道选项
            # self.pipeline_options = self.get_options() # Beam 2.x 推荐的方式
            # 或者通过main函数传递
            pass

        def process(self, element, otel_service_name, otel_resource_attributes):
            # 在这里使用 otel_service_name 和 otel_resource_attributes
            print(f"Processing with OTEL_SERVICE_NAME: {otel_service_name}")
            print(f"Processing with OTEL_RESOURCE_ATTRIBUTES: {otel_resource_attributes}")
            # 您的业务逻辑...
            yield element

    @staticmethod
    def parse_bill_data_requests(data):
        # 示例解析函数
        return data

def run_pipeline():
    # 从命令行或程序中解析管道选项
    # 注意：这里我们使用 CustomPipelineOptions 而不是 ProcessBillRequests.CustomOptions
    pipeline_options = CustomPipelineOptions()

    # 设置DataflowRunner所需的标准选项
    # 从环境变量获取或直接指定
    gcp_project_id = os.getenv("GCP_PROJECT_ID", "your-gcp-project")
    job_name = "process-bills-job"
    tas_gcs_bucket_name_prefix = os.getenv("TAS_GCS_BUCKET_NAME_PREFIX", "your-bucket-prefix")
    up_platform_env = os.getenv("UP_PLATFORM_ENV", "dev")
    service_account = os.getenv("SERVICE_ACCOUNT_EMAIL", "your-service-account@your-project.iam.gserviceaccount.com")
    subnetwork_url = os.getenv("SUBNETWORK_URL", None) # 例如 "regions/us-east1/subnetworks/default"
    uplight_telemetry_tar_file_path = "path/to/uplight-telemetry.tar.gz" # 替换为实际路径
    setup_file_path = "./setup.py" # 替换为实际路径

    # 将自定义选项的值传递给DataflowRunner的参数
    # DataflowRunner会从 pipeline_options 对象中解析这些值
    dataflow_options = pipeline_options.view_as(StandardOptions)
    dataflow_options.runner = 'DataflowRunner'
    dataflow_options.project = gcp_project_id
    dataflow_options.region = "us-east1"
    dataflow_options.job_name = job_name
    dataflow_options.temp_location = f'gs://{tas_gcs_bucket_name_prefix}{up_platform_env}/temp'
    dataflow_options.staging_location = f'gs://{tas_gcs_bucket_name_prefix}{up_platform_env}/staging'
    dataflow_options.save_main_session = True
    dataflow_options.service_account_email = service_account
    if subnetwork_url:
        dataflow_options.subnetwork = subnetwork_url
    dataflow_options.extra_packages = [uplight_telemetry_tar_file_path]
    dataflow_options.setup_file = setup_file_path

    # 获取自定义参数的值
    otel_service_name = pipeline_options.otel_service_name
    otel_resource_attributes = pipeline_options.otel_resource_attributes

    with beam.Pipeline(options=pipeline_options) as pipeline:
        read_from_db = beam.Create(["record1", "record2"]) # 模拟从DB读取

        result = (
            pipeline
            | "ReadPendingRecordsFromDB" >> read_from_db
            | "Parse input PCollection" >> beam.Map(ProcessBillRequests.parse_bill_data_requests)
            # 将自定义参数作为额外参数传递给DoFn
            | "Fetch bills " >> beam.ParDo(
                ProcessBillRequests.FetchBillInformation(),
                otel_service_name=otel_service_name,
                otel_resource_attributes=otel_resource_attributes
            )
        )
        pipeline.run().wait_until_finish()

if __name__ == '__main__':
    run_pipeline()

python your_pipeline_file.py \
    --runner=DataflowRunner \
    --project=your-gcp-project \
    --region=us-east1 \
    --job_name=my-bill-processing-job \
    --temp_location=gs://your-bucket/temp \
    --staging_location=gs://your-bucket/staging \
    --otel_service_name=my-billing-service \
    --otel_resource_attributes="env=prod,version=1.0"