通过Airflow插件发送PubSub通知后任务失败：psycopg2与SQLAlchemy报错

云上老周

2026-7-12

你的问题解答

是否有类似问题？
是的，不少Airflow用户在Composer环境的插件/钩子中使用同步IO的第三方客户端（包括PubSub）时遇到过类似的数据库连接异常，尤其是在任务状态变更的钩子中执行阻塞操作时。
psycopg2与PubSub客户端的兼容性问题？
psycopg2本身和PubSub客户端没有直接的API冲突，但存在两个关键的间接问题：

psycopg2的连接对象不支持多线程共享，PubSub客户端的后台线程可能意外访问了Airflow的数据库连接上下文，导致协议错误。
Composer环境中预安装的grpc、urllib3等底层依赖版本，可能同时被psycopg2和PubSub客户端依赖，版本不匹配时会引发底层网络/线程异常。

推荐的修复方案

方案1：异步处理PubSub发布，避免阻塞钩子

修改publish_event函数，不要同步等待future.result()，而是添加回调处理结果，让钩子快速完成，不阻塞Airflow的任务状态变更流程：

def publish_event(event: BaseEvent) -> None:
    """
    Publishes the given event to the configured Pub/Sub topic (异步非阻塞)
    """
    def publish_callback(future: pubsub_v1.publisher.futures.Future) -> None:
        try:
            response = future.result()
            logger.info(f"Published event: {event}. Response: {response}")
        except Exception as e:
            logger.error(f"Failed to publish event: {e}")

    try:
        logger.info(f"Attempting to publish event: {event}")
        future = publisher.publish(TOPIC, event.as_json_encoded())
        future.add_done_callback(publish_callback)
    except Exception as e:
        logger.error(f"Failed to initiate event publish: {e}")

方案2：将PubSub发布逻辑移至独立任务，而非钩子

放弃使用Airflow的状态钩子，改为在DAG中添加一个轻量级的任务（比如PythonOperator）来发送PubSub通知，这样可以隔离数据库连接和PubSub客户端的线程：

# 在你的DAG定义中添加
from airflow.operators.python import PythonOperator

def send_pubsub_notification(task_instance):
    # 复用现有handle_task_event逻辑
    handle_task_event(task_instance, "task_instance_success")

# 在目标任务完成后添加通知任务
dbt_task = DbtCloudRunJobOperator(...)
notify_task = PythonOperator(
    task_id="notify_pubsub",
    python_callable=send_pubsub_notification,
    op_kwargs={"task_instance": "{{ task_instance }}"},
    trigger_rule="all_done",  # 无论任务成功/失败都发送
)

dbt_task >> notify_task

方案3：调整PubSub客户端配置，隔离线程

初始化PubSub客户端时，指定独立的线程池，避免和Airflow的线程池共享：

from concurrent.futures import ThreadPoolExecutor

# 初始化publisher时指定独立线程池
publisher = pubsub_v1.PublisherClient(
    publisher_options=pubsub_v1.types.PublisherOptions(
        flow_control=pubsub_v1.types.PublishFlowControl(max_messages=1000),
    ),
    executor=ThreadPoolExecutor(max_workers=4),  # 独立的线程池
)

方案4：避免在钩子中修改数据库对象

你的handle_task_event中修改了dag_run.dag对象，这可能触发意外的数据库写入，建议改为只读访问：

def handle_task_event(task_instance: TaskInstance, event_type: str):
    dag_run = task_instance.dag_run
    # 改用只读方式获取dag信息，不修改dag_run对象
    dag = task_instance.task.dag or dag_run.dag
    dag_tags = dag.tags if dag else []
    
    dag_event = DagEvent(
        entity="task",
        event_type=event_type,
        dag_id=dag_run.dag_id,
        dag_status=dag_run.state,
        dag_start_date=dag_run.start_date,
        dag_end_date=dag_run.end_date,
        dag_tags=dag_tags,
    )
    task_event = TaskEvent(
        task_id=task_instance.task_id,
        task_status=task_instance.state,
        task_start_date=task_instance.start_date,
        task_end_date=task_instance.end_date,
        **dag_event.dict(),
    )
    publish_event(task_event)

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