如何通过配置conn_id实现PostgresOperator切换Airflow数据库环境

问题根源 #

PostgresOperator的postgres_conn_id默认不属于模板字段，直接传入模板字符串会被当作连接ID本身，而非解析后的值。你尝试自定义算子但未生效，核心原因是覆盖template_fields时丢失了父类的原有字段，导致模板解析逻辑未正确触发。

解决方案一：正确自定义PostgresOperator #

自定义算子时，需将postgres_conn_id追加到父类的template_fields中，而非直接覆盖：

from airflow.providers.postgres.operators.postgres import PostgresOperator
from typing import Sequence

class CustomPostgresOperator(PostgresOperator):
    # 继承父类原有模板字段并添加postgres_conn_id
    template_fields: Sequence[str] = (*PostgresOperator.template_fields, "postgres_conn_id")

之后在DAG中使用该自定义算子，你原有的模板字符串"{{ dag_run.conf.get('CONN_ID_TEST', 'pg_database') }}"就能被正确解析。

解决方案二：用PythonOperator结合PostgresHook动态执行SQL #

若不想自定义算子，可通过PythonOperator直接调用PostgresHook，完全依靠运行时配置控制连接：

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.operators.dummy import DummyOperator
import datetime as dt

DAG_ID = "init_database"

def run_postgres_query(**context):
    # 从运行时配置获取连接ID，默认使用生产库
    conn_id = context["dag_run"].conf.get("CONN_ID_TEST", "pg_database")
    hook = PostgresHook(postgres_conn_id=conn_id)
    # 执行目标SQL
    hook.run("SELECT * FROM pets LIMIT 1;", autocommit=True)

with DAG(
    dag_id=DAG_ID, 
    description="My dag",
    schedule_interval="@once",
    start_date=dt.datetime(2022, 1, 1),
    catchup=False,
) as dag:
    start = DummyOperator(task_id='start')
  
    my_task = PythonOperator(
        task_id="select",
        python_callable=run_postgres_query,
        provide_context=True  # 必须开启以获取dag_run等上下文变量
    )

start >> my_task

关键注意事项 #

• 连接环境变量命名修正：Airflow连接的环境变量格式为AIRFLOW_CONN_<连接ID大写>，你的变量名需调整为：
  • 生产库：AIRFLOW_CONN_PG_DATABASE（原AIRFLOW_CONN_ID_PG_DATABASE不符合规范）
  • 测试库：AIRFLOW_CONN_PG_DATABASE_TEST（原AIRFLOW_CONN_ID_PG_DATABASE_TEST不符合规范）
• 运行时配置验证：触发DAG时，确保传入的配置为{"CONN_ID_TEST": "pg_database_test"}，键名大小写需与代码保持一致。

内容的提问来源于stack exchange，提问作者Léo Guillaume