Fix VaultBackend.get_connection() breaking PythonVirtualenvOperator in Airflow 3

providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py

@@ -19,7 +19,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import json
 
 from airflow.providers.common.compat.sdk import conf
 from airflow.providers.hashicorp._internal_client.vault_client import _VaultClient
@@ -224,32 +224,33 @@ def _get_team_or_global_secret(self, base_path: str | None, team_name: str | Non
 
         return self._get_secret_with_base(path, key)
 
-    # Make sure connection is imported this way for type checking, otherwise when importing
-    # the backend it will get a circular dependency and fail
-    if TYPE_CHECKING:
-        from airflow.models.connection import Connection
-
-    def get_connection(self, conn_id: str, team_name: str | None = None) -> Connection | None:
+    def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | None:
         """
-        Get connection from Vault as secret.
+        Retrieve a connection from Vault as a serialized string.
 
-        Prioritize conn_uri if exists, if not fall back to normal Connection creation.
+        Returns the ``conn_uri`` value verbatim when present, otherwise serializes
+        the secret dict to JSON.  On Airflow 3.2+, the base-class ``get_connection``
+        deserializes the returned string using the Connection class that the framework
+        injects per execution context (ORM Connection on the server, SDK Connection in
+        workers), which avoids triggering SQLAlchemy mapper initialization in
+        task-execution subprocesses such as PythonVirtualenvOperator.  On the older
+        releases this provider still supports (2.11 / 3.0 / 3.1) there is no such
+        injection: the base ``deserialize_connection`` imports the ORM ``Connection``
+        directly, so the mapper-initialization avoidance does not apply there.
 
-        :return: A Connection object constructed from Vault data
+        :param conn_id: connection id
+        :param team_name: Team name associated to the task trying to access the connection (if any)
+        :return: Serialized connection string or None
         """
-        # The Connection needs to be locally imported because otherwise we get into cyclic import
-        # problems when instantiating the backend during configuration
-        from airflow.models.connection import Connection
-
         response = self._get_team_or_global_secret(self.connections_path, team_name, conn_id)
         if response is None:
             return None
 
         uri = response.get("conn_uri")
         if uri:
-            return Connection(conn_id, uri=uri)
+            return uri
 
-        return Connection(conn_id, **response)
+        return json.dumps(response)
 
     def get_variable(self, key: str, team_name: str | None = None) -> str | None:
         """

providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import json
 from unittest import mock
 
 import pytest
@@ -82,7 +83,7 @@ def variable_result(self):
         }
 
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
-    def test_get_connection(self, mock_hvac):
+    def test_get_conn_value(self, mock_hvac):
         mock_client = mock.MagicMock()
         mock_hvac.Client.return_value = mock_client
         mock_client.secrets.kv.v2.read_secret_version.return_value = {
@@ -121,8 +122,12 @@ def test_get_connection(self, mock_hvac):
         }
 
         test_client = VaultBackend(**kwargs)
-        connection = test_client.get_connection(conn_id="test_postgres")
-        assert connection.get_uri() == "postgresql://airflow:airflow@host:5432/airflow?foo=bar&baz=taz"
+        value = test_client.get_conn_value(conn_id="test_postgres")
+        assert value is not None
+        parsed = json.loads(value)
+        assert parsed["conn_type"] == "postgresql"
+        assert parsed["login"] == "airflow"
+        assert parsed["host"] == "host"
 
     @pytest.mark.parametrize(
         ("side_effects", "extra_kwargs", "exp_paths", "team_name"),
@@ -180,7 +185,7 @@ def test_get_connection_value_multi_team(
         )
 
         test_client = VaultBackend(**kwargs)
-        connection = test_client.get_connection(conn_id="test_postgres", team_name=team_name)
+        value = test_client.get_conn_value(conn_id="test_postgres", team_name=team_name)
         mock_client.secrets.kv.v2.read_secret_version.assert_has_calls(
             [
                 mock.call(
@@ -192,7 +197,8 @@ def test_get_connection_value_multi_team(
                 for path in exp_paths
             ]
         )
-        assert connection.get_uri() == "postgresql://airflow:airflow@host:5432/airflow?foo=bar&baz=taz"
+        assert value is not None
+        assert json.loads(value)["conn_type"] == "postgresql"
 
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
     def test_get_connection_without_predefined_mount_point(self, mock_hvac):
@@ -234,15 +240,16 @@ def test_get_connection_without_predefined_mount_point(self, mock_hvac):
         }
 
         test_client = VaultBackend(**kwargs)
-        connection = test_client.get_connection(conn_id="airflow/test_postgres")
-        assert connection.get_uri() == "postgresql://airflow:airflow@host:5432/airflow?foo=bar&baz=taz"
+        value = test_client.get_conn_value(conn_id="airflow/test_postgres")
+        assert value is not None
+        assert json.loads(value)["conn_type"] == "postgresql"
 
         # When mount_point=None and conn_id does not contain "/",
         # backend should return None and not call Vault
 
         mock_client.reset_mock()
 
-        assert test_client.get_connection("simple_id") is None
+        assert test_client.get_conn_value("simple_id") is None
         mock_client.secrets.kv.v2.read_secret_version.assert_not_called()
 
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
@@ -540,7 +547,7 @@ def test_auth_failure_raises_error(self, mock_hvac):
         }
 
         with pytest.raises(VaultError, match="Vault Authentication Error!"):
-            VaultBackend(**kwargs).get_connection(conn_id="test")
+            VaultBackend(**kwargs).get_conn_value(conn_id="test")
 
     def test_auth_type_kubernetes_with_unreadable_jwt_raises_error(self):
         path = "/var/tmp/this_does_not_exist/334e918ef11987d3ef2f9553458ea09f"
@@ -552,7 +559,7 @@ def test_auth_type_kubernetes_with_unreadable_jwt_raises_error(self):
         }
 
         with pytest.raises(FileNotFoundError, match=path):
-            VaultBackend(**kwargs).get_connection(conn_id="test")
+            VaultBackend(**kwargs).get_conn_value(conn_id="test")
 
     def test_auth_type_jwt_with_unreadable_jwt_raises_error(self):
         path = "/var/tmp/this_does_not_exist/jwt_token_file"
@@ -564,10 +571,10 @@ def test_auth_type_jwt_with_unreadable_jwt_raises_error(self):
         }
 
         with pytest.raises(FileNotFoundError, match=path):
-            VaultBackend(**kwargs).get_connection(conn_id="test")
+            VaultBackend(**kwargs).get_conn_value(conn_id="test")
 
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
-    def test_jwt_auth_type(self, mock_hvac):
+    def test_jwt_auth_type_conn_uri(self, mock_hvac):
         mock_client = mock.MagicMock()
         mock_hvac.Client.return_value = mock_client
         mock_client.secrets.kv.v2.read_secret_version.return_value = {
@@ -599,8 +606,55 @@ def test_jwt_auth_type(self, mock_hvac):
         }
 
         test_client = VaultBackend(**kwargs)
-        connection = test_client.get_connection(conn_id="test_postgres")
-        assert connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow"
+        value = test_client.get_conn_value(conn_id="test_postgres")
+        assert value == "postgresql://airflow:airflow@host:5432/airflow"
+        mock_client.auth.jwt.jwt_login.assert_called_with(
+            role="airflow-role", jwt="eyJhbGciOiJSUzI1NiJ9.test"
+        )
+
+    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
+    def test_jwt_auth_type_field_based(self, mock_hvac):
+        mock_client = mock.MagicMock()
+        mock_hvac.Client.return_value = mock_client
+        mock_client.secrets.kv.v2.read_secret_version.return_value = {
+            "request_id": "94011e25-f8dc-ec29-221b-1f9c1d9ad2ae",
+            "lease_id": "",
+            "renewable": False,
+            "lease_duration": 0,
+            "data": {
+                "data": {
+                    "conn_type": "postgres",
+                    "login": "airflow",
+                    "password": "airflow",
+                    "host": "host",
+                    "port": "5432",
+                    "schema": "airflow",
+                },
+                "metadata": {
+                    "created_time": "2020-03-16T21:01:43.331126Z",
+                    "deletion_time": "",
+                    "destroyed": False,
+                    "version": 1,
+                },
+            },
+            "wrap_info": None,
+            "warnings": None,
+            "auth": None,
+        }
+
+        kwargs = {
+            "connections_path": "connections",
+            "mount_point": "airflow",
+            "auth_type": "jwt",
+            "jwt_role": "airflow-role",
+            "jwt_token": "eyJhbGciOiJSUzI1NiJ9.test",
+            "url": "http://127.0.0.1:8200",
+        }
+
+        test_client = VaultBackend(**kwargs)
+        value = test_client.get_conn_value(conn_id="test_postgres")
+        assert value is not None
+        assert json.loads(value)["conn_type"] == "postgres"
         mock_client.auth.jwt.jwt_login.assert_called_with(
             role="airflow-role", jwt="eyJhbGciOiJSUzI1NiJ9.test"
         )
@@ -689,7 +743,7 @@ def test_connections_path_none_value(self, mock_hvac):
         }
 
         test_client = VaultBackend(**kwargs)
-        assert test_client.get_connection(conn_id="test") is None
+        assert test_client.get_conn_value(conn_id="test") is None
         mock_hvac.Client.assert_not_called()
 
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
@@ -699,7 +753,14 @@ def test_get_connection_with_empty_connections_path(self, mock_hvac):
 
         mock_client.secrets.kv.v2.read_secret_version.return_value = {
             "data": {
-                "data": {"conn_uri": "postgresql://user:pass@host:5432/db"},
+                "data": {
+                    "conn_type": "postgres",
+                    "login": "user",
+                    "password": "pass",
+                    "host": "host",
+                    "port": "5432",
+                    "schema": "db",
+                },
                 "metadata": {"version": 1},
             }
         }
@@ -714,7 +775,7 @@ def test_get_connection_with_empty_connections_path(self, mock_hvac):
 
         backend = VaultBackend(**kwargs)
 
-        connection = backend.get_connection("my_conn")
+        value = backend.get_conn_value("my_conn")
 
         # Assert Vault was called without "connections/" prefix
         mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with(
@@ -724,7 +785,11 @@ def test_get_connection_with_empty_connections_path(self, mock_hvac):
             raise_on_deleted_version=True,
         )
 
-        assert connection.get_uri() == "postgres://user:pass@host:5432/db"
+        assert value is not None
+        parsed = json.loads(value)
+        assert parsed["conn_type"] == "postgres"
+        assert parsed["login"] == "user"
+        assert parsed["host"] == "host"
 
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
     def test_variables_path_none_value(self, mock_hvac):