如何在Airflow DAG执行时为PostgresOperator配置动态数据库连接

使用Airflow的「Setup and Teardown」功能搭建临时PostgreSQL容器，流程如下：

1. 通过BashOperator启动PostGIS容器
2. 需要获取容器IP，用于配置PostgresOperator的连接
3. 尝试在DAG文件中直接读取BashOperator的XCOM值失败（DAG解析阶段无法获取运行时XCOM），且PostgresOperator无直接引用XCOM的参数

用户原代码：

from airflow.operators.postgres_operator import PostgresOperator
from airflow.operators.bash_operator import BashOperator
from airflow import DAG

from airflow.models.connection import Connection

# Define the DAG with the specified default_args
with DAG(
    "postgis_test_dag",
    schedule_interval=None,
) as dag:
    postgis_password = "postgis"
    postgis_user = "postgis"

    create_postgis = BashOperator(
        task_id="create_postgis",
        do_xcom_push=True,
        bash_command=f"docker run --name airflow_postgis -d postgis/postgis:17-3.5-alpine",
    )

    get_postgis_ip = BashOperator(
        task_id="get_postgis_ip",
        do_xcom_push=True,
        bash_command="docker inspect airflow_postgis -f '{{json.NetworkSettings.Networks}}'",
    )

    # This does not work
    postgis_ip = {{ TaskInstance.xcom_pull(task_ids='get_postgis_ip', key='return_value') }}

    #NEITHER:
    conn_id_postgis = f"postgresql://{postgis_user}:{postgis_password}@{postgis_ip}:5432/postgres"   
    os.environ['AIRFLOW_CONN_POSTGRES_DEFAULT'] = conn_id_postgis
    # NOR
    c = Connection(
        conn_id="postgis",
        host=postgis_ip,
        login="myname",
        port="port",
        password="mypassword")

    setup_postgis = PostgresOperator(
        task_id="setup_postgis",
        #conn_id=conn_id_postgis,
        sql="""SELECT * from network.roads(3381586.523108, 788516.801721)""",
    )

    delete_postgis = BashOperator(
        task_id="delete_postgis",
        bash_command="ID=$(docker container ls | grep postgis  | awk '{print $1}') && docker kill $ID  && docker rm $ID ",
    )

    create_postgis >> get_postgis_ip >> setup_postgis >> delete_postgis.as_teardown(setups=create_postgis)

Airflow的DAG分为解析阶段和运行阶段：

• 解析阶段：Airflow定期读取DAG文件，生成任务结构，此时无法获取运行时的XCOM值、任务实例（TaskInstance）等动态数据
• 运行阶段：任务实际执行，此时可以通过上下文（context）获取XCOM、任务实例等信息

你尝试在DAG解析阶段直接读取XCOM，这违反了Airflow的运行逻辑，因此必然失败。

方案1：通过PythonOperator动态创建Airflow Connection #

在运行阶段读取XCOM的容器IP，创建临时Connection供PostgresOperator使用：

from airflow.operators.postgres_operator import PostgresOperator
from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import PythonOperator
from airflow import DAG
from airflow.models.connection import Connection
from airflow.utils.session import provide_session

with DAG(
    "postgis_test_dag",
    schedule_interval=None,
) as dag:
    postgis_password = "postgis"
    postgis_user = "postgis"

    create_postgis = BashOperator(
        task_id="create_postgis",
        do_xcom_push=True,
        bash_command=f"docker run --name airflow_postgis -d postgis/postgis:17-3.5-alpine",
    )

    # 简化bash命令，直接提取容器IP（无需后续解析JSON）
    get_postgis_ip = BashOperator(
        task_id="get_postgis_ip",
        do_xcom_push=True,
        bash_command="docker inspect airflow_postgis -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}'",
    )

    @provide_session
    def create_temp_conn(**context, session=None):
        # 从上下文读取XCOM中的IP
        postgis_ip = context["ti"].xcom_pull(task_ids="get_postgis_ip", key="return_value")
        # 构建临时连接对象
        temp_conn = Connection(
            conn_id="postgis_temp",
            conn_type="postgres",
            host=postgis_ip,
            login=postgis_user,
            password=postgis_password,
            port=5432,
            schema="postgres"
        )
        # 清理同名旧连接，避免冲突
        existing_conn = session.query(Connection).filter(Connection.conn_id == temp_conn.conn_id).first()
        if existing_conn:
            session.delete(existing_conn)
        session.add(temp_conn)
        session.commit()

    setup_temp_conn = PythonOperator(
        task_id="setup_temp_conn",
        python_callable=create_temp_conn,
        provide_context=True,
    )

    setup_postgis = PostgresOperator(
        task_id="setup_postgis",
        conn_id="postgis_temp",  # 使用动态创建的临时连接
        sql="""SELECT * from network.roads(3381586.523108, 788516.801721)""",
    )

    delete_postgis = BashOperator(
        task_id="delete_postgis",
        bash_command="docker kill airflow_postgis && docker rm airflow_postgis",
    )

    # 任务依赖
    create_postgis >> get_postgis_ip >> setup_temp_conn >> setup_postgis >> delete_postgis.as_teardown(setups=create_postgis)

方案2：使用PostgresHook在PythonOperator中直接执行SQL #

无需创建Connection，直接在Python任务中读取XCOM的IP，用PostgresHook执行SQL：

from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import PythonOperator
from airflow import DAG
from airflow.hooks.postgres_hook import PostgresHook

with DAG(
    "postgis_test_dag",
    schedule_interval=None,
) as dag:
    postgis_password = "postgis"
    postgis_user = "postgis"

    create_postgis = BashOperator(
        task_id="create_postgis",
        do_xcom_push=True,
        bash_command=f"docker run --name airflow_postgis -d postgis/postgis:17-3.5-alpine",
    )

    get_postgis_ip = BashOperator(
        task_id="get_postgis_ip",
        do_xcom_push=True,
        bash_command="docker inspect airflow_postgis -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}'",
    )

    def run_postgis_query(**context):
        postgis_ip = context["ti"].xcom_pull(task_ids="get_postgis_ip", key="return_value")
        # 直接构建PostgresHook连接
        hook = PostgresHook(
            postgres_conn_id=None,
            host=postgis_ip,
            login=postgis_user,
            password=postgis_password,
            port=5432,
            schema="postgres"
        )
        sql = """SELECT * from network.roads(3381586.523108, 788516.801721)"""
        hook.run(sql)

    setup_postgis = PythonOperator(
        task_id="setup_postgis",
        python_callable=run_postgis_query,
        provide_context=True,
    )

    delete_postgis = BashOperator(
        task_id="delete_postgis",
        bash_command="docker kill airflow_postgis && docker rm airflow_postgis",
    )

    create_postgis >> get_postgis_ip >> setup_postgis >> delete_postgis.as_teardown(setups=create_postgis)

方案3：利用Docker容器DNS名称（简化方案） #

如果Airflow和PostGIS容器在同一Docker网络中，可以直接使用容器名称airflow_postgis作为主机名，无需获取IP：

from airflow.operators.postgres_operator import PostgresOperator
from airflow.operators.bash_operator import BashOperator
from airflow import DAG

with DAG(
    "postgis_test_dag",
    schedule_interval=None,
) as dag:
    postgis_password = "postgis"
    postgis_user = "postgis"

    create_postgis = BashOperator(
        task_id="create_postgis",
        # 确保容器加入Airflow所在的Docker网络（替换为实际网络名）
        bash_command=f"docker run --name airflow_postgis --network airflow_default -d postgis/postgis:17-3.5-alpine",
    )

    setup_postgis = PostgresOperator(
        task_id="setup_postgis",
        conn_id="postgis_default",  # 提前配置好使用容器名称的连接：postgresql://postgis:postgis@airflow_postgis:5432/postgres
        sql="""SELECT * from network.roads(3381586.523108, 788516.801721)""",
    )

    delete_postgis = BashOperator(
        task_id="delete_postgis",
        bash_command="docker kill airflow_postgis && docker rm airflow_postgis",
    )

    create_postgis >> setup_postgis >> delete_postgis.as_teardown(setups=create_postgis)

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