Unify connection not found exceptions between AF2 and AF3 #52968

amoghrajesh · 2025-07-07T08:37:30Z

Problem

Accessing a connection from SDK (Connection or using context) in airflow 3 and if connection isn't found it raises: AirflowRuntimeError: CONNECTION_NOT_FOUND: {'conn_id': 'my_conn'}

This is the case when pure SDK is used: airflow.sdk.Connection (airflow.models will work fine as we wrap it well)

Airflow 2 raises: airflow.exceptions.AirflowNotFoundException: The conn_id my_conn isn't defined .

Now due to this, some dags / providers use this for various sorts of filtering and for catching exceptions to perform some logic. So long term, this current way of handling exceptions is not feasible. We will start seeing problems when we start migrating providers to Airflow 3 entirely. Example issues: #52838 by amazon team and also #52921 by bosch team.

Testing the solution

DAG used:

from __future__ import annotations

from airflow.models.dag import dag
from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import Connection


def get_conn_context(**context):
    conn = context["conn"].my_conn
    return conn

def get_conn_sdk(**context):
    return Connection.get("my_conn")


def get_conn_models(**context):
    from airflow.models import Connection
    return Connection.get_connection_from_secrets("my_conn")

@dag()
def get_connection():

    a = PythonOperator(
        task_id="context_conn",
        python_callable=get_conn_context,
    )

    b = PythonOperator(
        task_id="sdk_conn",
        python_callable=get_conn_sdk,
    )

    c = PythonOperator(
        task_id="models_conn",
        python_callable=get_conn_models,
    )

    [a, b, c]


get_connection()

Without the fix:

Accessing from context(task a):

[2025-07-07, 13:41:04] INFO - DAG bundles loaded: dags-folder: source="airflow.dag_processing.bundles.manager.DagBundlesManager"
[2025-07-07, 13:41:04] INFO - Filling up the DagBag from /files/dags/get-connection.py: source="airflow.models.dagbag.DagBag"
[2025-07-07, 13:41:04] ERROR - Task failed with exception: source="task"
AirflowRuntimeError: CONNECTION_NOT_FOUND: {'conn_id': 'my_conn'}
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 875 in run

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 1162 in _execute_task

File "/opt/airflow/task-sdk/src/airflow/sdk/bases/operator.py", line 397 in wrapper

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 217 in execute

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 240 in execute_callable

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/callback_runner.py", line 82 in run

File "/files/dags/get-connection.py", line 9 in get_conn_context

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 275 in __getattr__

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 158 in _get_connection

Accessing from Models (task c):

[2025-07-07, 13:41:04] INFO - DAG bundles loaded: dags-folder: source="airflow.dag_processing.bundles.manager.DagBundlesManager"
[2025-07-07, 13:41:04] INFO - Filling up the DagBag from /files/dags/get-connection.py: source="airflow.models.dagbag.DagBag"
[2025-07-07, 13:41:04] ERROR - Task failed with exception: source="task"
AirflowNotFoundException: The conn_id `my_conn` isn't defined
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 875 in run

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 1162 in _execute_task

File "/opt/airflow/task-sdk/src/airflow/sdk/bases/operator.py", line 397 in wrapper

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 217 in execute

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 240 in execute_callable

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/callback_runner.py", line 82 in run

File "/files/dags/get-connection.py", line 18 in get_conn_models

File "/opt/airflow/airflow-core/src/airflow/models/connection.py", line 490 in get_connection_from_secrets

AirflowRuntimeError: CONNECTION_NOT_FOUND: {'conn_id': 'my_conn'}
File "/opt/airflow/airflow-core/src/airflow/models/connection.py", line 481 in get_connection_from_secrets

File "/opt/airflow/task-sdk/src/airflow/sdk/definitions/connection.py", line 152 in get

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 158 in _get_connection

Accessing from SDK:

[2025-07-07, 13:41:04] INFO - DAG bundles loaded: dags-folder: source="airflow.dag_processing.bundles.manager.DagBundlesManager"
[2025-07-07, 13:41:04] INFO - Filling up the DagBag from /files/dags/get-connection.py: source="airflow.models.dagbag.DagBag"
[2025-07-07, 13:41:04] ERROR - Task failed with exception: source="task"
AirflowRuntimeError: CONNECTION_NOT_FOUND: {'conn_id': 'my_conn'}
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 875 in run

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 1162 in _execute_task

File "/opt/airflow/task-sdk/src/airflow/sdk/bases/operator.py", line 397 in wrapper

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 217 in execute

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 240 in execute_callable

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/callback_runner.py", line 82 in run

File "/files/dags/get-connection.py", line 13 in get_conn_sdk

File "/opt/airflow/task-sdk/src/airflow/sdk/definitions/connection.py", line 152 in get

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 158 in _get_connection

After changes:

Accessing from context (task a):

[2025-07-07, 13:51:33] INFO - DAG bundles loaded: dags-folder: source="airflow.dag_processing.bundles.manager.DagBundlesManager"
[2025-07-07, 13:51:33] INFO - Filling up the DagBag from /files/dags/get-connection.py: source="airflow.models.dagbag.DagBag"
[2025-07-07, 13:51:33] ERROR - Task failed with exception: source="task"
AirflowNotFoundException: The conn_id `my_conn` isn't defined
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 875 in run

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 1162 in _execute_task

File "/opt/airflow/task-sdk/src/airflow/sdk/bases/operator.py", line 397 in wrapper

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 217 in execute

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 240 in execute_callable

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/callback_runner.py", line 82 in run

File "/files/dags/get-connection.py", line 9 in get_conn_context

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 280 in __getattr__

AirflowRuntimeError: CONNECTION_NOT_FOUND: {'conn_id': 'my_conn'}
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 277 in __getattr__

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 159 in _get_connection

Accessing from models (task b):

[2025-07-07, 13:51:33] INFO - DAG bundles loaded: dags-folder: source="airflow.dag_processing.bundles.manager.DagBundlesManager"
[2025-07-07, 13:51:33] INFO - Filling up the DagBag from /files/dags/get-connection.py: source="airflow.models.dagbag.DagBag"
[2025-07-07, 13:51:33] ERROR - Task failed with exception: source="task"
AirflowNotFoundException: The conn_id `my_conn` isn't defined
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 875 in run

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 1162 in _execute_task

File "/opt/airflow/task-sdk/src/airflow/sdk/bases/operator.py", line 397 in wrapper

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 217 in execute

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 240 in execute_callable

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/callback_runner.py", line 82 in run

File "/files/dags/get-connection.py", line 18 in get_conn_models

File "/opt/airflow/airflow-core/src/airflow/models/connection.py", line 481 in get_connection_from_secrets

File "/opt/airflow/task-sdk/src/airflow/sdk/definitions/connection.py", line 156 in get

AirflowRuntimeError: CONNECTION_NOT_FOUND: {'conn_id': 'my_conn'}
File "/opt/airflow/task-sdk/src/airflow/sdk/definitions/connection.py", line 153 in get

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 159 in _get_connection

Accessing from SDK (task c):

[2025-07-07, 13:51:33] INFO - DAG bundles loaded: dags-folder: source="airflow.dag_processing.bundles.manager.DagBundlesManager"
[2025-07-07, 13:51:33] INFO - Filling up the DagBag from /files/dags/get-connection.py: source="airflow.models.dagbag.DagBag"
[2025-07-07, 13:51:33] ERROR - Task failed with exception: source="task"
AirflowNotFoundException: The conn_id `my_conn` isn't defined
File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 875 in run

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/task_runner.py", line 1162 in _execute_task

File "/opt/airflow/task-sdk/src/airflow/sdk/bases/operator.py", line 397 in wrapper

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 217 in execute

File "/opt/airflow/providers/standard/src/airflow/providers/standard/operators/python.py", line 240 in execute_callable

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/callback_runner.py", line 82 in run

File "/files/dags/get-connection.py", line 13 in get_conn_sdk

File "/opt/airflow/task-sdk/src/airflow/sdk/definitions/connection.py", line 156 in get

AirflowRuntimeError: CONNECTION_NOT_FOUND: {'conn_id': 'my_conn'}
File "/opt/airflow/task-sdk/src/airflow/sdk/definitions/connection.py", line 153 in get

File "/opt/airflow/task-sdk/src/airflow/sdk/execution_time/context.py", line 159 in _get_connection

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

amoghrajesh · 2025-07-07T08:38:17Z

CC @jscheffl

ashb

We have the same try/except wrapping _get_connection in two places - is there a reason we don't handle inside that fn?

task-sdk/src/airflow/sdk/execution_time/context.py

ashb · 2025-07-07T09:05:50Z

Shouldn't this be a bug fix too?

amoghrajesh · 2025-07-07T09:54:39Z

We have the same try/except wrapping _get_connection in two places - is there a reason we don't handle inside that fn?

I checked on this and I see the general trend:

SDK client raised ErrorResponse with ErrorType for generic client related errors:

class ErrorType(enum.Enum):
    """Error types used in the API client."""

    CONNECTION_NOT_FOUND = "CONNECTION_NOT_FOUND"
    VARIABLE_NOT_FOUND = "VARIABLE_NOT_FOUND"
    XCOM_NOT_FOUND = "XCOM_NOT_FOUND"
    ASSET_NOT_FOUND = "ASSET_NOT_FOUND"
    DAGRUN_ALREADY_EXISTS = "DAGRUN_ALREADY_EXISTS"
    GENERIC_ERROR = "GENERIC_ERROR"
    API_SERVER_ERROR = "API_SERVER_ERROR"

These are captured and translated to AirflowRuntimeError in the task SDK runtime functions like connection, variables, dag run already exists etc. I guess this was done for convenience over using AirflowNotFoundException and other types of exceptions over just matching just AirflowRuntimeError?

amoghrajesh · 2025-07-07T09:55:17Z

Shouldn't this be a bug fix too?

Added 3.1 as I thought we are past the 3.0.3 deadline, but I realise we can do 3.0.4, marked it as such.

potiuk · 2025-07-07T10:11:09Z

Shouldn't this be a bug fix too?

Added 3.1 as I thought we are past the 3.0.3 deadline, but I realise we can do 3.0.4, marked it as such.

Also it's quite possible IMHO that we will 3.0.3rc4 - depending on @kaxil 's decision due to #52907

vincbeck · 2025-07-07T14:22:02Z

@ramitkataria , you might want to take a look a this one

task-sdk/src/airflow/sdk/execution_time/context.py

jscheffl

I can say this resolves the reported problem - Thanks!

ramitkataria

Looks good to me but I also agree with others about reducing code duplication where possible

ferruzzi

Also agree with it being a bugfix if we can still do that

kaxil · 2025-07-08T20:06:46Z

Also agree with it being a bugfix if we can still do that

yup, we can - 3.0.4 :)

amoghrajesh · 2025-07-09T06:48:52Z

@kaxil will merging this PR interfere with the RC-X of 3.0.3? If so, i can hold it. (Although we have the 304 milestone set)

kaxil · 2025-07-09T11:24:10Z

@kaxil will merging this PR interfere with the RC-X of 3.0.3? If so, i can hold it. (Although we have the 304 milestone set)

Go for it, I will workaround it :)

amoghrajesh · 2025-07-09T15:29:25Z

Cool thanks!

github-actions · 2025-07-09T15:30:34Z

Backport failed to create: v3-0-test. View the failure log Run details

Status	Branch	Result
❌	v3-0-test

You can attempt to backport this manually by running:

cherry_picker 5c2a2ec v3-0-test

This should apply the commit to the v3-0-test branch and leave the commit in conflict state marking
the files that need manual conflict resolution.

After you have resolved the conflicts, you can continue the backport process by running:

cherry_picker --continue

(cherry picked from commit 5c2a2ec)

amoghrajesh · 2025-07-09T15:41:09Z

Manual cherry pick here: #53093

…#52968) (#53093)

Unify connection not found exceptions between AF2 and AF3

564838f

amoghrajesh requested review from ashb and kaxil as code owners July 7, 2025 08:37

boring-cyborg bot added the area:task-sdk label Jul 7, 2025

amoghrajesh self-assigned this Jul 7, 2025

amoghrajesh requested review from ferruzzi, jscheffl, vatsrahul1001 and potiuk July 7, 2025 08:37

amoghrajesh added this to the Airflow 3.1.0 milestone Jul 7, 2025

ashb reviewed Jul 7, 2025

View reviewed changes

task-sdk/src/airflow/sdk/execution_time/context.py Outdated Show resolved Hide resolved

amoghrajesh modified the milestones: Airflow 3.1.0, Airflow 3.0.4 Jul 7, 2025

review comments from ash

01b7941

amoghrajesh requested a review from XD-DENG as a code owner July 7, 2025 10:01

amoghrajesh requested a review from ashb July 7, 2025 10:02

fixing tests

7806dd9

kaxil reviewed Jul 7, 2025

View reviewed changes

task-sdk/src/airflow/sdk/execution_time/context.py Outdated Show resolved Hide resolved

jscheffl approved these changes Jul 7, 2025

View reviewed changes

ramitkataria approved these changes Jul 8, 2025

View reviewed changes

review comments from kaxil

534147b

amoghrajesh requested a review from kaxil July 8, 2025 06:28

kaxil approved these changes Jul 8, 2025

View reviewed changes

ferruzzi approved these changes Jul 8, 2025

View reviewed changes

amoghrajesh mentioned this pull request Jul 9, 2025

Status of testing Providers that were prepared on July 08, 2025 #53016

Closed

17 tasks

kaxil added the backport-to-v3-0-test Mark PR with this label to backport to v3-0-test branch label Jul 9, 2025

amoghrajesh merged commit 5c2a2ec into apache:main Jul 9, 2025
78 checks passed

amoghrajesh deleted the better-connection-exceptions branch July 9, 2025 15:29

amoghrajesh added a commit to astronomer/airflow that referenced this pull request Jul 9, 2025

Unify connection not found exceptions between AF2 and AF3 (apache#52968)

586ba6e

(cherry picked from commit 5c2a2ec)

HsiuChuanHsu pushed a commit to HsiuChuanHsu/airflow that referenced this pull request Jul 10, 2025

Unify connection not found exceptions between AF2 and AF3 (apache#52968)

cfbc6af

amoghrajesh added a commit that referenced this pull request Jul 11, 2025

[v3-0-test] Unify connection not found exceptions between AF2 and AF3 (…

220fb39

…#52968) (#53093)

kaxil pushed a commit that referenced this pull request Jul 11, 2025

[v3-0-test] Unify connection not found exceptions between AF2 and AF3 (…

33cf422

…#52968) (#53093)

kaxil pushed a commit that referenced this pull request Jul 11, 2025

[v3-0-test] Unify connection not found exceptions between AF2 and AF3 (…

63c581e

…#52968) (#53093)

stephen-bracken pushed a commit to stephen-bracken/airflow that referenced this pull request Jul 15, 2025

Unify connection not found exceptions between AF2 and AF3 (apache#52968)

ddd27dd

Unify connection not found exceptions between AF2 and AF3 #52968

Unify connection not found exceptions between AF2 and AF3 #52968

Uh oh!

Conversation

amoghrajesh commented Jul 7, 2025 • edited by eladkal Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Testing the solution

Uh oh!

amoghrajesh commented Jul 7, 2025

Uh oh!

ashb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ashb commented Jul 7, 2025

Uh oh!

amoghrajesh commented Jul 7, 2025

Uh oh!

amoghrajesh commented Jul 7, 2025

Uh oh!

potiuk commented Jul 7, 2025

Uh oh!

vincbeck commented Jul 7, 2025

Uh oh!

Uh oh!

jscheffl left a comment

Choose a reason for hiding this comment

Uh oh!

ramitkataria left a comment

Choose a reason for hiding this comment

Uh oh!

ferruzzi left a comment

Choose a reason for hiding this comment

Uh oh!

kaxil commented Jul 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

amoghrajesh commented Jul 9, 2025

Uh oh!

kaxil commented Jul 9, 2025

Uh oh!

amoghrajesh commented Jul 9, 2025

Uh oh!

Uh oh!

github-actions bot commented Jul 9, 2025

Backport failed to create: v3-0-test. View the failure log Run details

Uh oh!

amoghrajesh commented Jul 9, 2025

Uh oh!

Uh oh!

amoghrajesh commented Jul 7, 2025 •

edited by eladkal

Loading

kaxil commented Jul 8, 2025 •

edited

Loading