Apache Beam PTransform 链式调用：构建高效数据处理管道

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions

# 1. 自定义 PTransform：从数据库读取数据
class ReadFromDatabase(beam.PTransform):
    """
    模拟从数据库读取数据的 PTransform。
    在实际应用中，会使用 beam.io.ReadFromJdbc 或其他数据库连接器。
    """
    def expand(self, pcoll):
        # 模拟从数据库读取的初始数据
        # 实际应从 pcoll 参数接收一个 PCollection，但这里为了演示，
        # 我们从无到有创建一个 PCollection。
        # 在实际管道中，pcoll 可能是 pipeline 对象，例如：
        # return pcoll | 'ReadRecords' >> beam.io.ReadFromJdbc(...)
        print("--- Step 1: Reading from Database ---")
        return pcoll | 'CreateInitialRecords' >> beam.Create([
            {'id': 1, 'name': 'ProductA', 'category': 'Electronics'},
            {'id': 2, 'name': 'ProductB', 'category': 'Books'}
        ])

# 2. 自定义 DoFn：调用第一个 REST API
class CallFirstAPI(beam.DoFn):
    """
    对每个数据库记录调用第一个外部 REST API。
    假设 API 返回一个包含子项的数组。
    """
    def process(self, element):
        # 模拟 API 调用逻辑
        print(f"--- Step 2: Calling First API for ID: {element['id']} ---")
        # 假设 API 返回一个包含子项的数组
        if element['id'] == 1:
            api_response_array = [{'sub_id': 'A1', 'value': 100}, {'sub_id': 'A2', 'value': 150}]
        else:
            api_response_array = [{'sub_id': 'B1', 'value': 200}]

        # 将原始数据与 API 响应结合，传递给下一步
        yield {
            'id': element['id'],
            'name': element['name'],
            'category': element['category'],
            'first_api_data': api_response_array # 包含数组的响应
        }

# 3. 自定义 DoFn：处理 API 响应数组并调用第二个 API
class ProcessArrayAndCallSecondAPI(beam.DoFn):
    """
    接收包含数组的 PCollection 元素，对数组中的每个子项调用第二个 API，
    并产生新的 PCollection 元素（扁平化处理）。
    """
    def process(self, element):
        record_id = element['id']
        first_api_data_array = element['first_api_data']

        print(f"--- Step 3: Processing Array and Calling Second API for ID: {record_id} ---")
        for sub_item in first_api_data_array:
            # 模拟调用第二个 API
            # 假设第二个 API 返回一些补充信息
            second_api_info = f"info_for_{sub_item['sub_id']}"

            # 组合所有相关数据，作为新的元素输出
            yield {
                'id': record_id,
                'name': element['name'],
                'category': element['category'],
                'sub_id': sub_item['sub_id'],
                'value': sub_item['value'],
                'second_api_info': second_api_info
            }

# 4. 自定义 DoFn：更新数据到数据库
class UpdateDatabase(beam.DoFn):
    """
    模拟将最终处理结果更新到数据库。
    在实际应用中，会使用 beam.io.WriteToJdbc 或自定义的数据库写入逻辑。
    """
    def process(self, element):
        # 模拟数据库更新操作
        print(f"--- Step 4: Updating Database for ID: {element['id']}, Sub_ID: {element['sub_id']} with data: {element} ---")
        # 实际中会执行 INSERT/UPDATE 语句
        # 例如：db_connection.execute("UPDATE ... WHERE id = ? AND sub_id = ?", element['id'], element['sub_id'])
        yield element # 可以选择不返回，或者返回更新成功的标识

# 构建 Beam 管道
def run_pipeline():
    with beam.Pipeline(options=PipelineOptions()) as pipeline:
        # Step 1: 从数据库读取初始记录
        # 注意：这里 ReadFromDatabase 接收 pipeline 对象作为输入，
        # 因为它负责创建初始的 PCollection。
        initial_records = pipeline | 'ReadFromDB' >> ReadFromDatabase()

        # Step 2: 对每条记录调用第一个 API
        first_api_results = initial_records | 'CallFirstAPI' >> beam.ParDo(CallFirstAPI())

        # Step 3: 处理第一个 API 的响应数组，并调用第二个 API
        # 注意：这里使用 ParDo(DoFn) 来实现扁平化和多步处理
        final_processed_data = first_api_results | 'ProcessArrayAndCallSecondAPI' >> beam.ParDo(ProcessArrayAndCallSecondAPI())

        # Step 4: 将最终处理结果更新到数据库
        # 这里可以使用 beam.Map 打印最终结果，或者用 beam.io.WriteToJdbc
        final_processed_data | 'LogFinalResults' >> beam.Map(print)
        # 实际的数据库更新步骤
        # final_processed_data | 'UpdateDB' >> beam.ParDo(UpdateDatabase())

if __name__ == '__main__':
    run_pipeline()