Python实现Apache Beam WriteToDatastore结果写入BigQuery审计表

我有一段简单的Python Apache Beam管道代码：

with beam.Pipeline(options=create_pipeline_options(pipeline_args)) as p:
            rows = (p | 'ReadFromBigquery' >> beam.io.ReadFromBigQuery(table=f"{known_args.project}:{known_args.datasetId}.{known_args.tableId}", use_standard_sql=True))
            entities  = (rows | 'GetEntities' >> beam.ParDo(GetEntity()))
            updated = (entities | 'Update Entities' >> beam.ParDo(UpdateEntity()))
            _ = (updated | 'Write To Datastore' >> WriteToDatastore(known_args.project))

我希望在WriteToDatastore执行完成后，记录哪些实体已成功更新，并将其写入BigQuery审计表，理想实现如下：

successful_entities, failed entities = (updated | 'Write To Datastore' >> WriteToDatastoreWrapper(known_args.project))
 _ = (successful_entities | 'Write Success To Bigquery' >> beam.io.WriteToBigQuery(table=f"{c.audit_table}:{known_args.datasetId}.{known_args.tableId}"))
 _ = (failed_entities| 'Write Failed To Bigquery' >> beam.io.WriteToBigQuery(table=f"{c.audit_table}:{known_args.datasetId}.{known_args.tableId}"))

请问该需求是否可实现？另外，若批量数据经过n次重试后仍失败，能否结合runId捕获失败并记录对应的失败批次？

1. 需求完全可实现 #

要实现分流成功/失败实体的需求，有两种可行方案：

方案一：自定义DoFn包装Datastore写入逻辑 #

自己实现一个WriteToDatastoreWrapper类，在其中执行Datastore写入操作，捕获单个实体的异常，将成功和失败结果通过TaggedOutput分流到不同的PCollection。示例代码如下：

import datetime
from google.cloud import datastore
import apache_beam as beam

class WriteToDatastoreWrapper(beam.DoFn):
    def __init__(self, project):
        self.project = project
        self.datastore_client = None

    def setup(self):
        # 初始化Datastore客户端，仅在worker启动时执行一次
        self.datastore_client = datastore.Client(project=self.project)

    def process(self, entity):
        try:
            self.datastore_client.put(entity)
            # 输出成功记录，包含实体ID、状态、时间戳
            yield beam.pvalue.TaggedOutput(
                'success', 
                {'entity_id': entity.key.id(), 'status': 'success', 'timestamp': datetime.datetime.now()}
            )
        except Exception as e:
            # 输出失败记录，附加错误信息
            yield beam.pvalue.TaggedOutput(
                'failure', 
                {'entity_id': entity.key.id(), 'status': 'failure', 'error': str(e), 'timestamp': datetime.datetime.now()}
            )

# 管道中使用方式
with beam.Pipeline(options=create_pipeline_options(pipeline_args)) as p:
    rows = (p | 'ReadFromBigquery' >> beam.io.ReadFromBigQuery(
        table=f"{known_args.project}:{known_args.datasetId}.{known_args.tableId}", 
        use_standard_sql=True
    ))
    entities = (rows | 'GetEntities' >> beam.ParDo(GetEntity()))
    updated = (entities | 'Update Entities' >> beam.ParDo(UpdateEntity()))
    
    # 执行写入并分流结果
    results = (updated 
               | 'Write To Datastore' >> beam.ParDo(WriteToDatastoreWrapper(known_args.project))
               .with_outputs('success', 'failure'))
    
    # 写入成功审计表
    _ = (results.success | 'Write Success To Bigquery' >> beam.io.WriteToBigQuery(
        table=f"{known_args.project}:{known_args.datasetId}.audit_success",
        schema='entity_id:STRING, status:STRING, timestamp:TIMESTAMP',
        write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND
    ))
    
    # 写入失败审计表
    _ = (results.failure | 'Write Failed To Bigquery' >> beam.io.WriteToBigQuery(
        table=f"{known_args.project}:{known_args.datasetId}.audit_failure",
        schema='entity_id:STRING, status:STRING, error:STRING, timestamp:TIMESTAMP',
        write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND
    ))

方案二：使用Beam内置Try转换简化分流 #

Beam的Try转换可以自动将成功处理的元素和抛出异常的元素分离，无需手动捕获异常：

from apache_beam.utils.try_fn import Try
from google.cloud import datastore

def write_to_datastore(entity, project):
    client = datastore.Client(project=project)
    client.put(entity)
    return {'entity_id': entity.key.id(), 'status': 'success', 'timestamp': datetime.datetime.now()}

with beam.Pipeline(options=create_pipeline_options(pipeline_args)) as p:
    # ... 前面的读取、转换逻辑省略
    # 执行写入并分流
    attempt_results = updated | 'Attempt Datastore Write' >> Try(write_to_datastore, known_args.project)
    
    # 拆分成功/失败结果
    successful_entities = attempt_results | 'Filter Success' >> beam.Filter(lambda x: x.success)
    failed_entities = attempt_results | 'Filter Failure' >> beam.Filter(lambda x: not x.success)
    
    # 转换失败结果格式（保留错误信息）
    formatted_failed = failed_entities | 'Format Failed' >> beam.Map(
        lambda x: {'entity_id': x.exception.args[0].key.id(), 'status': 'failure', 'error': str(x.exception), 'timestamp': datetime.datetime.now()}
    )
    
    # 写入审计表逻辑同方案一

2. 结合RunID捕获重试后失败的批次 #

完全可以实现，步骤如下：

步骤1：获取管道RunID #

RunID是Beam管道运行的唯一标识，可通过管道选项获取：

from apache_beam.options.pipeline_options import RunOptions

pipeline_options = create_pipeline_options(pipeline_args)
run_id = pipeline_options.view_as(RunOptions).run_id

步骤2：实现带重试的写入逻辑 #

使用重试库（如tenacity）实现n次重试逻辑，重试耗尽仍失败时，将RunID、批次标识（可选）写入审计表：

from tenacity import retry, stop_after_attempt, retry_if_exception_type
from google.cloud import datastore

class RetryingDatastoreWriter(beam.DoFn):
    def __init__(self, project, run_id, max_retries=3):
        self.project = project
        self.run_id = run_id
        self.max_retries = max_retries
        self.datastore_client = None

    def setup(self):
        self.datastore_client = datastore.Client(project=self.project)

    # 定义重试规则：仅针对Datastore相关异常重试，最多3次
    @retry(
        stop=stop_after_attempt(3),
        retry=retry_if_exception_type(datastore.exceptions.DatastoreError)
    )
    def _attempt_write(self, entity):
        self.datastore_client.put(entity)

    def process(self, entity, batch_id=None):
        try:
            self._attempt_write(entity)
            yield beam.pvalue.TaggedOutput('success', {
                'entity_id': entity.key.id(),
                'run_id': self.run_id,
                'batch_id': batch_id,
                'status': 'success',
                'timestamp': datetime.datetime.now()
            })
        except Exception as e:
            yield beam.pvalue.TaggedOutput('failure', {
                'entity_id': entity.key.id(),
                'run_id': self.run_id,
                'batch_id': batch_id,
                'status': 'failure',
                'error': str(e),
                'retry_count': self.max_retries,
                'timestamp': datetime.datetime.now()
            })

# 管道中使用
with beam.Pipeline(options=pipeline_options) as p:
    # ... 前面的读取、转换逻辑省略
    # 可选：给实体添加批次ID（按每1000个实体为一批）
    updated_with_batch = (updated 
                          | 'Add Batch ID' >> beam.Map(
                              lambda elem, idx: (elem, idx//1000), 
                              beam.utilities.core.globally()
                          ))
    
    # 执行带重试的写入并分流
    results = (updated_with_batch 
               | 'Write With Retry' >> beam.ParDo(RetryingDatastoreWriter(known_args.project, run_id))
               .with_outputs('success', 'failure'))
    
    # 写入审计表逻辑同方案一

注意事项 #

• 审计表建议分开成功/失败表，或用单表加status字段区分，便于后续查询分析。
• Datastore写入需保证幂等性，避免重试导致重复写入数据。
• 本地运行时RunID由job_name生成，生产环境Dataflow会自动分配唯一RunID。

内容的提问来源于stack exchange，提问作者Giuseppe Joe Ramundo