如何实现基于PubSub通知触发的Dataflow BigQuery读取任务？

方案1：流式Dataflow持续监听PubSub #

这种方式适合需要长期监听PubSub Topic，每次收到符合条件的消息就执行BigQuery读取的场景。

完整代码示例

import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
import json
from google.cloud import bigquery

def parse_pubsub_message(message):
    """解析PubSub消息为JSON对象"""
    try:
        return json.loads(message.decode('utf-8'))
    except json.JSONDecodeError:
        return None

def filter_valid_triggers(msg):
    """过滤出符合条件的触发消息"""
    return msg is not None and msg.get('job_id') == 101 and msg.get('status') == 'Success'

def fetch_bigquery_data(_):
    """读取BigQuery数据并返回结果列表"""
    query = '''
        SELECT * FROM `bigquery-public-data.chicago_crime.crime` LIMIT 100
    '''
    client = bigquery.Client()
    query_job = client.query(query)
    return [dict(row) for row in query_job.result()]

if __name__ == '__main__':
    # 配置Dataflow运行参数
    pipeline_options = PipelineOptions(
        runner='DataflowRunner',
        project='your-project-id',
        region='your-region',
        staging_location='gs://your-bucket/staging',
        temp_location='gs://your-bucket/temp',
        streaming=True  # 启用流式模式持续监听PubSub
    )

    with beam.Pipeline(options=pipeline_options) as p:
        (
            p
            # 从目标PubSub Topic读取消息
            | '读取PubSub消息' >> beam.io.ReadFromPubSub(topic='projects/your-project-id/topics/your-topic')
            # 解析消息为JSON
            | '解析JSON' >> beam.Map(parse_pubsub_message)
            # 过滤仅保留符合条件的触发消息
            | '过滤有效触发' >> beam.Filter(filter_valid_triggers)
            # 触发BigQuery读取，展开查询结果
            | '读取BigQuery数据' >> beam.FlatMap(fetch_bigquery_data)
            # 输出结果（可替换为你的数据转换逻辑）
            | '打印结果' >> beam.Map(print)
        )

关键注意事项

• 权限配置：确保Dataflow服务账号拥有PubSub订阅权限、BigQuery数据读取权限以及GCS存储桶的读写权限。
• 重复触发处理：如果同一触发消息可能多次发送，可通过在消息中加入唯一ID，结合Beam的状态管理（beam.StateSpec）记录已处理的任务，避免重复执行。
• 资源消耗：流式Dataflow会持续运行，需根据消息频率评估资源配置。

方案2：Cloud Functions触发批处理Dataflow #

如果仅需要在收到符合条件的消息时启动一次性的BigQuery读取任务，这种方案更节省资源，成本更低。

Cloud Functions代码示例

import json
from googleapiclient.discovery import build

def trigger_dataflow_job(event, context):
    """监听PubSub消息，触发符合条件的Dataflow批处理任务"""
    pubsub_message = json.loads(event['data'].decode('utf-8'))
    
    # 校验触发条件
    if pubsub_message.get('job_id') != 101 or pubsub_message.get('status') != 'Success':
        return
    
    # 构建Dataflow API请求
    dataflow_client = build('dataflow', 'v1b3')
    project_id = 'your-project-id'
    job_config = {
        'jobName': f'bigquery-crime-read-{context.event_id}',
        'environment': {'tempLocation': 'gs://your-bucket/temp'},
        'steps': [
            {
                'name': 'ReadFromBigQuery',
                'bigquery': {
                    'query': '''SELECT * FROM `bigquery-public-data.chicago_crime.crime` LIMIT 100''',
                    'useStandardSql': True
                }
            },
            {
                'name': 'ProcessData',
                'map': {
                    'language': 'PYTHON',
                    'code': 'def process(element):\n    print(element)\n    return element'
                }
            }
        ]
    }
    
    # 启动Dataflow任务
    request = dataflow_client.projects().locations().jobs().create(
        projectId=project_id,
        location='your-region',
        body=job_config
    )
    response = request.execute()
    print(f"已触发Dataflow任务：{response['id']}")

方案说明

1. 创建Cloud Function并订阅目标PubSub Topic。
2. 当收到消息时，先校验job_id和status，符合条件则调用Dataflow API启动批处理任务。
3. Dataflow任务执行完成后自动终止，无需持续占用资源。

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