Apache Beam PTransform输出传递与复杂数据流构建实践

import apache_beam as beam
import requests # 实际API调用可能用到

# 1. 模拟从数据库读取数据的PTransform
class ReadFromDatabase(beam.PTransform):
    def expand(self, pcoll):
        # 在实际应用中，这里会使用 beam.io.ReadFromJdbc 或其他数据库连接器
        # 模拟读取两行数据，每行是一个字典
        print("Executing ReadFromDatabase...")
        return pcoll | 'ReadDatabaseEntries' >> beam.Create([
            {'id': 1, 'name': 'Alice', 'email': 'alice@example.com'},
            {'id': 2, 'name': 'Bob', 'email': 'bob@example.com'}
        ])

# 2. 调用第一个REST API的PTransform
class CallFirstAPI(beam.PTransform):
    class ProcessElement(beam.DoFn):
        def process(self, element):
            # 模拟调用第一个API，并获取一个包含数组的响应
            # 实际中会使用 requests.get(f"http://api.example.com/first/{element['id']}")
            print(f"CallFirstAPI - Processing element: {element['id']}")
            api_response = {
                'status': 'success',
                'data': {
                    'id': element['id'],
                    'details': f"details_for_{element['name']}",
                    'items': [f"itemA_{element['id']}", f"itemB_{element['id']}"] # 模拟数组
                }
            }
            # 将原始数据与API响应合并，并传递给下一步
            yield {**element, 'first_api_data': api_response['data']}

    def expand(self, pcoll):
        return pcoll | 'CallFirstAPIProcess' >> beam.ParDo(self.ProcessElement())

# 3. 调用第二个REST API的PTransform (针对数组中的每个元素)
class CallSecondAPI(beam.PTransform):
    class ProcessElement(beam.DoFn):
        def process(self, element):
            first_api_data = element['first_api_data']
            items = first_api_data.get('items', [])

            # 对第一个API响应中的每个item调用第二个API
            for item in items:
                # 模拟调用第二个API
                # 实际中会使用 requests.get(f"http://api.example.com/second/{item}")
                print(f"CallSecondAPI - Processing item: {item} for element: {element['id']}")
                second_api_response = {
                    'item_name': item,
                    'additional_info': f"info_for_{item}"
                }
                # 将原始数据、第一个API数据和当前第二个API响应合并
                # 注意：这里可能会产生多个输出元素，每个对应一个item
                yield {
                    **element,
                    'current_item_data': second_api_response
                }

    def expand(self, pcoll):
        # 使用ParDo处理每个元素，并可能产生多个输出
        return pcoll | 'CallSecondAPIProcess' >> beam.ParDo(self.ProcessElement())

# 4. 模拟更新数据库的PTransform (仅作示意)
class UpdateDatabase(beam.PTransform):
    class ProcessElement(beam.DoFn):
        def process(self, element):
            # 在实际应用中，这里会使用 beam.io.WriteToJdbc 或其他数据库写入器
            # 可能需要根据element中的id和API数据构建SQL UPDATE语句
            print(f"UpdateDatabase - Updating record: {element['id']} with data: {element}")
            # 实际场景中，此DoFn可能不会yield任何元素，或者yield一个更新成功标记
            yield element # 仅仅为了在管道末尾查看数据流

    def expand(self, pcoll):
        return pcoll | 'UpdateDatabaseEntries' >> beam.ParDo(self.ProcessElement())

# 构建并运行Beam管道
with beam.Pipeline() as pipeline:
    # 步骤1: 从数据库读取数据
    read_from_db_pcoll = pipeline | 'Start' >> ReadFromDatabase()

    # 步骤2: 调用第一个API，其输入是 read_from_db_pcoll 的输出
    call_first_api_pcoll = read_from_db_pcoll | 'CallFirstAPI' >> CallFirstAPI()

    # 步骤3: 调用第二个API，其输入是 call_first_api_pcoll 的输出
    # 注意：CallSecondAPI可能会将一个输入元素扩展为多个输出元素
    call_second_api_pcoll = call_first_api_pcoll | 'CallSecondAPI' >> CallSecondAPI()

    # 步骤4: 更新数据库，其输入是 call_second_api_pcoll 的输出
    # 在实际场景中，可能需要进一步聚合或处理 call_second_api_pcoll 的输出
    # 例如，如果需要将多个item的结果聚合回原始记录，可能需要使用GroupByKey
    call_second_api_pcoll | 'UpdateDatabase' >> UpdateDatabase()

    # 如果需要查看最终结果，可以写入文件或打印
    # call_second_api_pcoll | 'WriteToConsole' >> beam.Map(print)