Prioritize secrets backend over DB for retrieving connections (apache#47593)

Author: amoghrajesh

airflow/hooks/base.py

@@ -20,7 +20,6 @@
 from __future__ import annotations
 
 import logging
-import sys
 from typing import TYPE_CHECKING, Any, Protocol
 
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -60,18 +59,6 @@ def get_connection(cls, conn_id: str) -> Connection:
         :param conn_id: connection id
         :return: connection
         """
-        # TODO: This is not the best way of having compat, but it's "better than erroring" for now. This still
-        # means SQLA etc is loaded, but we can't avoid that unless/until we add import shims as a big
-        # back-compat layer
-
-        # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
-        # and should use the Task SDK API server path
-        if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
-            # TODO: AIP 72: Add deprecation here once we move this module to task sdk.
-            from airflow.sdk import Connection as TaskSDKConnection
-
-            return TaskSDKConnection.get(conn_id=conn_id)
-
         from airflow.models.connection import Connection
 
         conn = Connection.get_connection_from_secrets(conn_id)

airflow/models/connection.py

@@ -20,6 +20,7 @@
 import json
 import logging
 import re
+import sys
 from contextlib import suppress
 from json import JSONDecodeError
 from typing import Any
@@ -34,6 +35,7 @@
 from airflow.models.crypto import get_fernet
 from airflow.sdk.execution_time.secrets_masker import mask_secret
 from airflow.secrets.cache import SecretCache
+from airflow.secrets.metastore import MetastoreBackend
 from airflow.utils.helpers import prune_dict
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.module_loading import import_string
@@ -446,6 +448,8 @@ def get_connection_from_secrets(cls, conn_id: str) -> Connection:
         """
         Get connection by conn_id.
 
+        If `MetastoreBackend` is getting used in the execution context, use Task SDK API.
+
         :param conn_id: connection id
         :return: connection
         """
@@ -459,6 +463,18 @@ def get_connection_from_secrets(cls, conn_id: str) -> Connection:
 
         # iterate over backends if not in cache (or expired)
         for secrets_backend in ensure_secrets_loaded():
+            if isinstance(secrets_backend, MetastoreBackend):
+                # TODO: This is not the best way of having compat, but it's "better than erroring" for now. This still
+                # means SQLA etc is loaded, but we can't avoid that unless/until we add import shims as a big
+                # back-compat layer
+
+                # If this is set it means are in some kind of execution context (Task, Dag Parse or Triggerer perhaps)
+                # and should use the Task SDK API server path
+                if hasattr(sys.modules.get("airflow.sdk.execution_time.task_runner"), "SUPERVISOR_COMMS"):
+                    # TODO: AIP 72: Add deprecation here once we move this module to task sdk.
+                    from airflow.sdk import Connection as TaskSDKConnection
+
+                    return TaskSDKConnection.get(conn_id=conn_id)
             try:
                 conn = secrets_backend.get_connection(conn_id=conn_id)
                 if conn:

tests/hooks/test_base.py

@@ -17,7 +17,22 @@
 # under the License.
 from __future__ import annotations
 
+from unittest import mock
+
+import pytest
+
 from airflow.hooks.base import BaseHook
+from airflow.sdk.execution_time.comms import ConnectionResult, GetConnection
+
+from tests_common.test_utils.config import conf_vars
+
+
+@pytest.fixture
+def mock_supervisor_comms():
+    with mock.patch(
+        "airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS", create=True
+    ) as supervisor_comms:
+        yield supervisor_comms
 
 
 class TestBaseHook:
@@ -32,3 +47,42 @@ def test_custom_logger_name_is_correctly_set(self):
     def test_empty_string_as_logger_name(self):
         hook = BaseHook(logger_name="")
         assert hook.log.name == "airflow.task.hooks"
+
+    def test_get_connection(self, mock_supervisor_comms):
+        conn = ConnectionResult(
+            conn_id="test_conn",
+            conn_type="mysql",
+            host="mysql",
+            schema="airflow",
+            login="root",
+            password="password",
+            port=1234,
+            extra='{"extra_key": "extra_value"}',
+        )
+
+        mock_supervisor_comms.get_message.return_value = conn
+
+        hook = BaseHook(logger_name="")
+        hook.get_connection(conn_id="test_conn")
+
+        mock_supervisor_comms.send_request.assert_called_once_with(
+            msg=GetConnection(conn_id="test_conn"), log=mock.ANY
+        )
+
+    def test_get_connection_secrets_backend_configured(self, mock_supervisor_comms, tmp_path):
+        path = tmp_path / "conn.env"
+        path.write_text("CONN_A=mysql://host_a")
+
+        with conf_vars(
+            {
+                ("secrets", "backend"): "airflow.secrets.local_filesystem.LocalFilesystemBackend",
+                ("secrets", "backend_kwargs"): f'{{"connections_file_path": "{path}"}}',
+            }
+        ):
+            hook = BaseHook(logger_name="")
+            retrieved_conn = hook.get_connection(conn_id="CONN_A")
+
+            assert retrieved_conn.conn_id == "CONN_A"
+
+            mock_supervisor_comms.send_request.assert_not_called()
+            mock_supervisor_comms.get_message.assert_not_called()