在 Apache Beam 上按顺序执行数据库写入和读取任务

Question

我正在使用 beam-nuggets 库从 Apache Beam 管道读取数据并将数据写入 Postgres 数据库。我想依次执行以下两个任务：

1. 在processing_info 表中插入新行。
2. 查询新创建的processing_info表记录的主键。然后，主键查询的值将作为侧输入传递给另一个DoFn（用于填充相关表的外键列）。​​

目前，我在执行 Beam 管道之前创建 processing_info 记录，但我想创建新记录作为管道执行的一部分（在 Google Dataflow 上运行管道时，它使事情变得更简单）。理想情况下，代码应如下所示：

with beam.Pipeline(options=pipeline_options) as p:
    # Executes first
    proc_id_result = (p | 'Create Proc Info Record' >> beam.Create([{'pipeline_name': 'cleansed_data_pipeline'}])
                          | 'Make Processing Id' >> relational_db.Write(
                                source_config=source_config,
                                table_config=proc_table_config))
    # Executes second
    proc_id_record = p | relational_db.ReadFromDB(
            source_config=source_config,
            table_name='processing_info',
            query='SELECT pi.id FROM processing_info pi WHERE processing_date_time = '
                  '   (SELECT MAX(pi1.processing_date_time) from processing_info pi1 '
                  f'      where pi1.pipeline_name = \'cleansed_data_pipeline\')'
        )
    ...
    # This code executes later, and is automatically deferred until the side input is available
    | 'Add \'processing_info_id\'' >>
                (beam.ParDo(AddKeyValuePairToDict(), 'processing_info_id', AsSingleton(proc_id_record)))
    ...

我可能能够破解某些内容（例如未使用的侧输入）以推迟查询直到插入操作完成，但我想知道是否有更惯用的方法（我是 Beam 的新手）。

谢谢。

Answer 1

您想到了正确的想法：您可以使用未使用的侧输入来做到这一点。你会做这样的事情（在 Beam 本身中用于 ReadFromBigQuery

class PassThrough(beam.DoFn):
  def process(self, element):
    yield element

output = input | beam.ParDo(PassThrough()).with_outputs(
    'cleanup_signal', main='main')
main_output = output['main']
cleanup_signal = output['cleanup_signal']

single_element = (
    input.pipeline
    | beam.Create([None])
    | beam.Map(lambda x, nothing: x, beam.pvalue.AsSingleton(cleanup_signal)))

single_element | relational_db.ReadFromDB(...)

现在，问题在于使用您的 ReadFromDB 转换来实现这项工作，我猜它不会接受这样的输入。有没有办法实现这种转换？