Airflow PubSubPullSensor/Operator：下游任务如何确认PubSub消息

核心方案 #

Airflow没有专门的内置Operator用于单独确认PubSub消息，需通过PythonOperator调用PubSub API，在下游任务成功后完成消息确认。

分步实现 #

1. 配置消息拉取任务
   使用PubSubPullOperator或PubSubPullSensor时，务必将ack_messages设为False，此时消息不会被自动确认，需由下游任务处理确认逻辑。官方文档说明：

   若ack_messages设为True，消息会在返回前立即被确认；否则，下游任务需负责确认消息。
   同时开启do_xcom_push=True，将拉取到的消息（包含ack_id）通过XCom传递给下游任务。

2. 编写消息确认函数
   基于Google Cloud PubSub SDK编写确认逻辑，示例代码如下：

   from google.cloud import pubsub_v1
   
   def ack_pubsub_messages(project_id: str, subscription_id: str, ack_ids: list):
       subscriber = pubsub_v1.SubscriberClient()
       subscription_path = subscriber.subscription_path(project_id, subscription_id)
       # 批量确认消息
       subscriber.acknowledge(
           request={"subscription": subscription_path, "ack_ids": ack_ids}
       )
       subscriber.close()
   

3. 构建任务依赖链
   确保确认任务仅在下游处理任务成功后执行，示例DAG结构：

   pull_task = PubSubPullOperator(
       task_id="pull_pubsub_msgs",
       project_id="your-gcp-project",
       subscription_id="your-subscription",
       ack_messages=False,
       do_xcom_push=True
   )
   
   process_task = PythonOperator(
       task_id="process_msgs",
       python_callable=your_process_logic,
       op_args=[{{ ti.xcom_pull(task_ids="pull_pubsub_msgs") }}]
   )
   
   ack_task = PythonOperator(
       task_id="ack_pubsub_msgs",
       python_callable=ack_pubsub_messages,
       op_args=[
           "your-gcp-project",
           "your-subscription",
           {{ ti.xcom_pull(task_ids="pull_pubsub_msgs") | map(attribute='ack_id') | list }}
       ]
   )
   
   pull_task >> process_task >> ack_task
   

关键注意事项 #

• 确保Airflow运行环境具备PubSub订阅的pubsub.subscriptions.acknowledge权限，可通过服务账号密钥配置。
• 若下游处理任务失败，消息会按照订阅的重试策略重新进入待拉取队列，无需额外处理。

内容的提问来源于stack exchange，提问作者Despicable me