无法更新自定义模板的Dataflow批处理作业参数

问题根源 #

1. 模板创建时固化业务参数：你创建模板时传入了--db、--dbCSKey等业务参数，这些值会被直接固化到模板中，运行时无法覆盖。
2. 构建阶段执行MongoDB读取：代码在Pipeline构建（模板创建）阶段就直接连接MongoDB并执行collection.find()，导致模板中包含了创建时的查询结果，运行时不会重新读取新数据源。
3. 未使用动态参数机制：直接读取my_options.db这类属性会在构建阶段求值，无法在运行时获取启动作业时传入的新参数。

修复步骤 #

1. 修改代码，实现运行时动态参数读取与延迟数据源加载

将MongoDB读取逻辑延迟到作业运行时执行，使用RuntimeValueProvider动态获取参数：

import apache_beam as beam
from apache_beam.io.gcp.bigquery import WriteToBigQuery
from apache_beam.options.pipeline_options import PipelineOptions, SetupOptions, RuntimeValueProvider
import logging
import json
from google.cloud import secretmanager
import pymongo

logging.basicConfig(level=logging.INFO)
  
class MyOptions(PipelineOptions):
    @classmethod
    def _add_argparse_args(cls, parser):
        parser.add_argument('--db', required=True)
        parser.add_argument('--dbCSKey', required=True)
        parser.add_argument('--collectionName', required=True)
        parser.add_argument('--clientId', required=True)
        parser.add_argument('--projectId', required=True)
        parser.add_argument('--bigQuerySchema', required=True)
        parser.add_argument('--secretName', required=True)

# 自定义DoFn，在运行时读取MongoDB数据
class ReadMongoDB(beam.DoFn):
    def process(self, _, db, dbCSKey, collectionName, secretName):
        # 从SecretManager获取MongoDB连接串
        client = secretmanager.SecretManagerServiceClient()
        response = client.access_secret_version(request={'name': secretName})
        secrets = response.payload.data.decode('UTF-8')
        secretDict = json.loads(secrets)
        mongoURI = next((tenant['connectionString'] for tenant in secretDict['tenants'] if tenant['clientId'] == dbCSKey), None)
        
        # 连接MongoDB并读取数据
        client = pymongo.MongoClient(mongoURI)
        db_instance = client[db]
        collection = db_instance[collectionName]
        yield from collection.find()

pipeline_options = PipelineOptions()
my_options = pipeline_options.view_as(MyOptions)
# 确保worker能加载主会话依赖
pipeline_options.view_as(SetupOptions).save_main_session = True
    
with beam.Pipeline(options=pipeline_options) as p:
    # 从RuntimeValueProvider获取运行时参数
    db = RuntimeValueProvider.get_option('db', str, None)
    dbCSKey = RuntimeValueProvider.get_option('dbCSKey', str, None)
    collectionName = RuntimeValueProvider.get_option('collectionName', str, None)
    secretName = RuntimeValueProvider.get_option('secretName', str, None)
    projectId = RuntimeValueProvider.get_option('projectId', str, None)
    bigQuerySchema = RuntimeValueProvider.get_option('bigQuerySchema', str, None)
    
    bigQueryTable = f'{projectId}:{bigQuerySchema}.{collectionName}'
    
    # 通过空PCollection触发MongoDB读取，传递运行时参数
    mongo_data = (
        p 
        | '启动读取流程' >> beam.Create([None])
        | '读取MongoDB数据' >> beam.ParDo(ReadMongoDB(), db=db, dbCSKey=dbCSKey, collectionName=collectionName, secretName=secretName)
    )
    
    # 写入BigQuery的逻辑（补充完整你的业务流程）
    mongo_data | '写入BigQuery' >> WriteToBigQuery(
        bigQueryTable,
        write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
        create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED
    )

2. 修改模板创建命令，移除业务参数

创建模板时仅传入模板构建必需的参数，业务参数留到启动作业时传入：

python3 -m data-sync --region us-west1 --runner DataflowRunner --project datalake-dev-358216 --temp_location gs://datalake-dev-358216.appspot.com/temp --template_location gs://datalake-dev-358216.appspot.com/template

3. 启动模板时传入业务参数

保持原启动命令逻辑（可去掉重复的job_name参数，因为命令开头的名称即为作业名）：

gcloud dataflow jobs run data-sync-initial-load-humana-insurance-claim-procedures-20230612040559 --gcs-location gs://datalake-dev-358216.appspot.com/template --parameters db=insurance_010_dev,dbCSKey=humana,collectionName=insurance_claim_procedures,clientId=humana,projectId=datalake-dev-358216,bigQuerySchema=ml_data_marts,secretName=projects/762809992414/secrets/mongo-multi-tenant-secret/versions/latest

关键说明 #

• RuntimeValueProvider：确保参数在作业运行时动态获取，避免模板构建阶段固化值。
• 延迟数据源读取：将MongoDB连接、查询逻辑放到ParDo的process方法中，确保这部分逻辑仅在作业运行时执行。
• 模板与业务参数分离：模板创建时只传递构建所需参数，业务参数在启动作业时传入，保证模板的通用性。

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