From abfdd80120d467743b269ac7a2d2c78454afbe66 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Tue, 9 Jun 2026 13:51:28 -0500 Subject: [PATCH 1/9] Fix VaultBackend.get_connection() causing mapper init failure in virtualenv tasks Constructing airflow.models.connection.Connection (SQLAlchemy ORM) triggers lazy mapper initialisation for the entire Airflow model registry. In PythonVirtualenvOperator subprocesses, DagModel has not been imported so DagScheduleAssetNameReference cannot resolve its 'DagModel' relationship, raising sqlalchemy.exc.InvalidRequestError. This exception is silently swallowed in context.py's secrets-backend loop, making the connection appear undefined even when Vault returns a valid 200 response. Switch to airflow.providers.common.compat.sdk.Connection (Pydantic in Airflow 3, SQLAlchemy in Airflow 2) which does not trigger mapper init. Also update call sites to use keyword-only conn_id= to match the SDK Connection's constructor signature. --- .../providers/hashicorp/secrets/vault.py | 28 ++++-- .../unit/hashicorp/secrets/test_vault.py | 88 ++++++++++++++++++- 2 files changed, 108 insertions(+), 8 deletions(-) diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index 72ac30901b2e3..384151bd4d554 100644 --- a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py +++ b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py @@ -227,7 +227,7 @@ def _get_team_or_global_secret(self, base_path: str | None, team_name: str | Non # 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 + from airflow.providers.common.compat.sdk import Connection def get_connection(self, conn_id: str, team_name: str | None = None) -> Connection | None: """ @@ -237,9 +237,10 @@ def get_connection(self, conn_id: str, team_name: str | None = None) -> Connecti :return: A Connection object constructed from Vault data """ - # 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 + # Use the compat SDK Connection (Pydantic in Airflow 3, SQLAlchemy in Airflow 2) to avoid + # triggering SQLAlchemy mapper initialization for unrelated models (e.g. DagModel) in + # task-execution subprocesses such as PythonVirtualenvOperator. + from airflow.providers.common.compat.sdk import Connection response = self._get_team_or_global_secret(self.connections_path, team_name, conn_id) if response is None: @@ -247,9 +248,24 @@ def get_connection(self, conn_id: str, team_name: str | None = None) -> Connecti uri = response.get("conn_uri") if uri: - return Connection(conn_id, uri=uri) + # from_uri is available on the SDK Connection (Airflow 3.2+ / task-sdk 1.2.0+). + # On Airflow 2 the SQLAlchemy Connection accepts uri= in its __init__. + # On Airflow 3.0/3.1 the attrs-based Connection has neither from_uri nor a uri= + # constructor arg; return None so the secrets-backend loop tries the next backend. + if hasattr(Connection, "from_uri"): + return Connection.from_uri(uri, conn_id=conn_id) # type: ignore[attr-defined] + try: + return Connection(conn_id=conn_id, uri=uri) # type: ignore[call-arg] + except TypeError: + self.log.warning( + "Cannot deserialize conn_uri for connection '%s': upgrade to Airflow 3.2+ " + "or store the connection using individual fields (conn_type, host, login, " + "etc.) in Vault.", + conn_id, + ) + return None - return Connection(conn_id, **response) + return Connection(conn_id=conn_id, **response) def get_variable(self, key: str, team_name: str | None = None) -> str | None: """ diff --git a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py index 8b6af85a725da..ba762520a7b71 100644 --- a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py +++ b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py @@ -21,6 +21,7 @@ import pytest from hvac.exceptions import InvalidPath, VaultError +from airflow.providers.common.compat.sdk import Connection as SdkConnection from airflow.providers.hashicorp.secrets.vault import VaultBackend from tests_common.test_utils.config import conf_vars @@ -576,7 +577,14 @@ def test_jwt_auth_type(self, mock_hvac): "renewable": False, "lease_duration": 0, "data": { - "data": {"conn_uri": "postgresql://airflow:airflow@host:5432/airflow"}, + "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": "", @@ -699,7 +707,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}, } } @@ -759,3 +774,72 @@ def test_config_path_none_value(self, mock_hvac): test_client = VaultBackend(**kwargs) assert test_client.get_config("test") is None mock_hvac.Client.assert_not_called() + + @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") + def test_get_connection_returns_sdk_connection_not_sqlalchemy_model(self, mock_hvac): + """get_connection must return the SDK (Pydantic) Connection, not the SQLAlchemy ORM model. + + Using the SQLAlchemy model triggers lazy mapper initialisation for the entire Airflow model + registry, which fails in PythonVirtualenvOperator subprocesses where DagModel has not been + imported (raises sqlalchemy.exc.InvalidRequestError, silently swallowed in context.py). + """ + mock_client = mock.MagicMock() + mock_hvac.Client.return_value = mock_client + mock_client.secrets.kv.v2.read_secret_version.return_value = { + "data": { + "data": { + "conn_type": "postgres", + "host": "db-host", + "login": "user", + "password": "pass", + "port": "5432", + "schema": "mydb", + }, + "metadata": {"version": 1}, + } + } + + backend = VaultBackend( + connections_path="connections", + mount_point="airflow", + auth_type="token", + url="http://127.0.0.1:8200", + token="token", + ) + conn = backend.get_connection("trino_default") + + assert isinstance(conn, SdkConnection), ( + f"Expected SDK Connection, got {type(conn)}. " + "Returning the SQLAlchemy model triggers mapper init and breaks PythonVirtualenvOperator." + ) + assert conn.conn_id == "trino_default" + + @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") + @pytest.mark.skipif( + not hasattr(SdkConnection, "from_uri"), + reason="conn_uri deserialization requires Connection.from_uri (Airflow 3.2+ / task-sdk 1.2.0+)", + ) + def test_get_connection_via_uri_returns_sdk_connection(self, mock_hvac): + mock_client = mock.MagicMock() + mock_hvac.Client.return_value = mock_client + mock_client.secrets.kv.v2.read_secret_version.return_value = { + "data": { + "data": {"conn_uri": "postgresql://user:pass@host:5432/db"}, + "metadata": {"version": 1}, + } + } + + backend = VaultBackend( + connections_path="connections", + mount_point="airflow", + auth_type="token", + url="http://127.0.0.1:8200", + token="token", + ) + conn = backend.get_connection("my_conn") + + assert isinstance(conn, SdkConnection), ( + f"Expected SDK Connection, got {type(conn)}. " + "Returning the SQLAlchemy model triggers mapper init and breaks PythonVirtualenvOperator." + ) + assert conn.conn_id == "my_conn" From 53250e37424a70df05c4ee7326b47d1323192a95 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Wed, 10 Jun 2026 10:51:22 -0500 Subject: [PATCH 2/9] Remove unnecessary type: ignore comments on Connection.from_uri and uri= calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit from_uri is a defined classmethod on airflow.sdk.Connection, and the uri= overload is visible via the manual __init__ overloads — neither suppression is needed. --- .../src/airflow/providers/hashicorp/secrets/vault.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index 384151bd4d554..928e1f9cad52a 100644 --- a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py +++ b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py @@ -253,9 +253,9 @@ def get_connection(self, conn_id: str, team_name: str | None = None) -> Connecti # On Airflow 3.0/3.1 the attrs-based Connection has neither from_uri nor a uri= # constructor arg; return None so the secrets-backend loop tries the next backend. if hasattr(Connection, "from_uri"): - return Connection.from_uri(uri, conn_id=conn_id) # type: ignore[attr-defined] + return Connection.from_uri(uri, conn_id=conn_id) try: - return Connection(conn_id=conn_id, uri=uri) # type: ignore[call-arg] + return Connection(conn_id=conn_id, uri=uri) except TypeError: self.log.warning( "Cannot deserialize conn_uri for connection '%s': upgrade to Airflow 3.2+ " From 8e682f96c4395e791855ba00ba4837c1c328a8d4 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Wed, 10 Jun 2026 11:57:10 -0500 Subject: [PATCH 3/9] Widen except clause to catch ValueError in addition to TypeError Pydantic or future attrs validators can raise ValueError; catching only TypeError would let those slip through silently the same way this PR fixes the original silent failure. --- .../hashicorp/src/airflow/providers/hashicorp/secrets/vault.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index 928e1f9cad52a..db94ccba58e30 100644 --- a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py +++ b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py @@ -256,7 +256,7 @@ def get_connection(self, conn_id: str, team_name: str | None = None) -> Connecti return Connection.from_uri(uri, conn_id=conn_id) try: return Connection(conn_id=conn_id, uri=uri) - except TypeError: + except (TypeError, ValueError): self.log.warning( "Cannot deserialize conn_uri for connection '%s': upgrade to Airflow 3.2+ " "or store the connection using individual fields (conn_type, host, login, " From e03b429db46d185c55aa4ffe8b7f3fb277a00dfc Mon Sep 17 00:00:00 2001 From: seanmuth Date: Fri, 12 Jun 2026 11:55:41 -0500 Subject: [PATCH 4/9] Restore type: ignore[call-arg] on Connection(uri=) fallback The attrs mypy plugin overrides the manually-written __init__ overloads, so mypy sees only the attrs-generated __init__ (no uri= param). The ignore is required on the Airflow 2 / pre-3.2 fallback path. --- .../hashicorp/src/airflow/providers/hashicorp/secrets/vault.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index db94ccba58e30..8c90fff9639ce 100644 --- a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py +++ b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py @@ -255,7 +255,7 @@ def get_connection(self, conn_id: str, team_name: str | None = None) -> Connecti if hasattr(Connection, "from_uri"): return Connection.from_uri(uri, conn_id=conn_id) try: - return Connection(conn_id=conn_id, uri=uri) + return Connection(conn_id=conn_id, uri=uri) # type: ignore[call-arg] except (TypeError, ValueError): self.log.warning( "Cannot deserialize conn_uri for connection '%s': upgrade to Airflow 3.2+ " From 0560c15101374bba6cb55a9d32e558b3e5fbb1ec Mon Sep 17 00:00:00 2001 From: seanmuth Date: Fri, 26 Jun 2026 10:20:32 -0500 Subject: [PATCH 5/9] Use get_conn_value() and framework-injected Connection class MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the get_connection() override (which hard-coded the compat SDK Connection) with get_conn_value(), returning a URI string for conn_uri secrets or a JSON-serialized field dict otherwise. The base-class get_connection() then deserializes using _get_connection_class(), which the framework populates with the ORM Connection on the server and the SDK Connection in workers — fixing the mapper-init regression without breaking server-side commands like airflow connections get. Tests call _set_connection_class(SdkConnection) to simulate the framework injection, matching how initialize_secrets_backends() wires things up at runtime. --- .../providers/hashicorp/secrets/vault.py | 46 ++++++------------- .../unit/hashicorp/secrets/test_vault.py | 13 ++++-- 2 files changed, 24 insertions(+), 35 deletions(-) diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index 8c90fff9639ce..85fbb5dfad0e9 100644 --- a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py +++ b/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,48 +224,30 @@ 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.providers.common.compat.sdk 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. The base-class ``get_connection`` deserializes the + returned string using the Connection class that the framework injected for the + current execution context (ORM Connection on the server, SDK Connection in + workers), which avoids triggering SQLAlchemy mapper initialization in + task-execution subprocesses such as PythonVirtualenvOperator. - :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 """ - # Use the compat SDK Connection (Pydantic in Airflow 3, SQLAlchemy in Airflow 2) to avoid - # triggering SQLAlchemy mapper initialization for unrelated models (e.g. DagModel) in - # task-execution subprocesses such as PythonVirtualenvOperator. - from airflow.providers.common.compat.sdk 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: - # from_uri is available on the SDK Connection (Airflow 3.2+ / task-sdk 1.2.0+). - # On Airflow 2 the SQLAlchemy Connection accepts uri= in its __init__. - # On Airflow 3.0/3.1 the attrs-based Connection has neither from_uri nor a uri= - # constructor arg; return None so the secrets-backend loop tries the next backend. - if hasattr(Connection, "from_uri"): - return Connection.from_uri(uri, conn_id=conn_id) - try: - return Connection(conn_id=conn_id, uri=uri) # type: ignore[call-arg] - except (TypeError, ValueError): - self.log.warning( - "Cannot deserialize conn_uri for connection '%s': upgrade to Airflow 3.2+ " - "or store the connection using individual fields (conn_type, host, login, " - "etc.) in Vault.", - conn_id, - ) - return None + return uri - return Connection(conn_id=conn_id, **response) + return json.dumps(response) def get_variable(self, key: str, team_name: str | None = None) -> str | None: """ diff --git a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py index ba762520a7b71..4398b400b1218 100644 --- a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py +++ b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py @@ -122,8 +122,9 @@ def test_get_connection(self, mock_hvac): } test_client = VaultBackend(**kwargs) + test_client._set_connection_class(SdkConnection) connection = test_client.get_connection(conn_id="test_postgres") - assert connection.get_uri() == "postgresql://airflow:airflow@host:5432/airflow?foo=bar&baz=taz" + assert connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow?foo=bar&baz=taz" @pytest.mark.parametrize( ("side_effects", "extra_kwargs", "exp_paths", "team_name"), @@ -181,6 +182,7 @@ def test_get_connection_value_multi_team( ) test_client = VaultBackend(**kwargs) + test_client._set_connection_class(SdkConnection) connection = test_client.get_connection(conn_id="test_postgres", team_name=team_name) mock_client.secrets.kv.v2.read_secret_version.assert_has_calls( [ @@ -193,7 +195,7 @@ 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 connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow?foo=bar&baz=taz" @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_get_connection_without_predefined_mount_point(self, mock_hvac): @@ -235,8 +237,9 @@ def test_get_connection_without_predefined_mount_point(self, mock_hvac): } test_client = VaultBackend(**kwargs) + test_client._set_connection_class(SdkConnection) connection = test_client.get_connection(conn_id="airflow/test_postgres") - assert connection.get_uri() == "postgresql://airflow:airflow@host:5432/airflow?foo=bar&baz=taz" + assert connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow?foo=bar&baz=taz" # When mount_point=None and conn_id does not contain "/", # backend should return None and not call Vault @@ -607,6 +610,7 @@ def test_jwt_auth_type(self, mock_hvac): } test_client = VaultBackend(**kwargs) + test_client._set_connection_class(SdkConnection) connection = test_client.get_connection(conn_id="test_postgres") assert connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow" mock_client.auth.jwt.jwt_login.assert_called_with( @@ -728,6 +732,7 @@ def test_get_connection_with_empty_connections_path(self, mock_hvac): } backend = VaultBackend(**kwargs) + backend._set_connection_class(SdkConnection) connection = backend.get_connection("my_conn") @@ -806,6 +811,7 @@ def test_get_connection_returns_sdk_connection_not_sqlalchemy_model(self, mock_h url="http://127.0.0.1:8200", token="token", ) + backend._set_connection_class(SdkConnection) conn = backend.get_connection("trino_default") assert isinstance(conn, SdkConnection), ( @@ -836,6 +842,7 @@ def test_get_connection_via_uri_returns_sdk_connection(self, mock_hvac): url="http://127.0.0.1:8200", token="token", ) + backend._set_connection_class(SdkConnection) conn = backend.get_connection("my_conn") assert isinstance(conn, SdkConnection), ( From d35144ce8b419cdeb237b388d387e2957d6184ea Mon Sep 17 00:00:00 2001 From: seanmuth Date: Fri, 26 Jun 2026 11:03:39 -0500 Subject: [PATCH 6/9] Guard _set_connection_class calls for Airflow 3.0/2.11 compat _set_connection_class was added after 3.0/2.11; guard test calls with hasattr and skip the two isinstance(SdkConnection) tests on older Airflow where class injection is unavailable. --- .../unit/hashicorp/secrets/test_vault.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py index 4398b400b1218..4a2848587fc41 100644 --- a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py +++ b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py @@ -122,7 +122,8 @@ def test_get_connection(self, mock_hvac): } test_client = VaultBackend(**kwargs) - test_client._set_connection_class(SdkConnection) + if hasattr(test_client, "_set_connection_class"): + test_client._set_connection_class(SdkConnection) connection = test_client.get_connection(conn_id="test_postgres") assert connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow?foo=bar&baz=taz" @@ -182,7 +183,8 @@ def test_get_connection_value_multi_team( ) test_client = VaultBackend(**kwargs) - test_client._set_connection_class(SdkConnection) + if hasattr(test_client, "_set_connection_class"): + test_client._set_connection_class(SdkConnection) connection = test_client.get_connection(conn_id="test_postgres", team_name=team_name) mock_client.secrets.kv.v2.read_secret_version.assert_has_calls( [ @@ -237,7 +239,8 @@ def test_get_connection_without_predefined_mount_point(self, mock_hvac): } test_client = VaultBackend(**kwargs) - test_client._set_connection_class(SdkConnection) + if hasattr(test_client, "_set_connection_class"): + test_client._set_connection_class(SdkConnection) connection = test_client.get_connection(conn_id="airflow/test_postgres") assert connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow?foo=bar&baz=taz" @@ -610,7 +613,8 @@ def test_jwt_auth_type(self, mock_hvac): } test_client = VaultBackend(**kwargs) - test_client._set_connection_class(SdkConnection) + if hasattr(test_client, "_set_connection_class"): + test_client._set_connection_class(SdkConnection) connection = test_client.get_connection(conn_id="test_postgres") assert connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow" mock_client.auth.jwt.jwt_login.assert_called_with( @@ -732,7 +736,8 @@ def test_get_connection_with_empty_connections_path(self, mock_hvac): } backend = VaultBackend(**kwargs) - backend._set_connection_class(SdkConnection) + if hasattr(backend, "_set_connection_class"): + backend._set_connection_class(SdkConnection) connection = backend.get_connection("my_conn") @@ -781,6 +786,10 @@ def test_config_path_none_value(self, mock_hvac): mock_hvac.Client.assert_not_called() @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") + @pytest.mark.skipif( + not hasattr(VaultBackend, "_set_connection_class"), + reason="Connection class injection requires BaseSecretsBackend._set_connection_class (Airflow 3.2+)", + ) def test_get_connection_returns_sdk_connection_not_sqlalchemy_model(self, mock_hvac): """get_connection must return the SDK (Pydantic) Connection, not the SQLAlchemy ORM model. @@ -822,8 +831,8 @@ def test_get_connection_returns_sdk_connection_not_sqlalchemy_model(self, mock_h @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") @pytest.mark.skipif( - not hasattr(SdkConnection, "from_uri"), - reason="conn_uri deserialization requires Connection.from_uri (Airflow 3.2+ / task-sdk 1.2.0+)", + not hasattr(SdkConnection, "from_uri") or not hasattr(VaultBackend, "_set_connection_class"), + reason="conn_uri deserialization requires Connection.from_uri and _set_connection_class (Airflow 3.2+)", ) def test_get_connection_via_uri_returns_sdk_connection(self, mock_hvac): mock_client = mock.MagicMock() From 06d07c36ad623358be420cf5e602a1cf7a92b2f4 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Fri, 26 Jun 2026 12:08:22 -0500 Subject: [PATCH 7/9] Re-add get_connection() override to accept team_name on all Airflow versions On Airflow 3.0/2.11, BaseSecretsBackend.get_connection() did not accept team_name. Without a VaultBackend override, calling get_connection with team_name raised TypeError. The override delegates to get_conn_value() then uses deserialize_connection() when available (modern Airflow, class injection works), and falls back to the compat SDK Connection on older Airflow that predates _get_connection_class / deserialize_connection. --- .../providers/hashicorp/secrets/vault.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index 85fbb5dfad0e9..8d83e17bec87c 100644 --- a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py +++ b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py @@ -224,6 +224,44 @@ def _get_team_or_global_secret(self, base_path: str | None, team_name: str | Non return self._get_secret_with_base(path, key) + def get_connection(self, conn_id: str, team_name: str | None = None): + """ + Return connection object with a given ``conn_id``. + + Overrides the base to accept ``team_name`` on all supported Airflow versions and + to fall back to the compat SDK Connection when the framework-level class injection + (``deserialize_connection`` / ``_get_connection_class``) is not yet available. + + :param conn_id: connection id + :param team_name: Team name associated to the task trying to access the connection (if any) + :return: Connection object or None + """ + value = self.get_conn_value(conn_id=conn_id, team_name=team_name) + if value is None: + return None + # Prefer the base-class deserializer when available — it uses _get_connection_class() + # which the framework populates with the right Connection class per execution context. + if hasattr(self, "deserialize_connection"): + return self.deserialize_connection(conn_id=conn_id, value=value) + # Fallback for older Airflow that predates _get_connection_class / deserialize_connection. + from airflow.providers.common.compat.sdk import Connection + + value = value.strip() + if value.startswith("{"): + return Connection.from_json(value=value, conn_id=conn_id) + if hasattr(Connection, "from_uri"): + return Connection.from_uri(uri=value, conn_id=conn_id) + try: + return Connection(conn_id=conn_id, uri=value) # type: ignore[call-arg] + except (TypeError, ValueError): + self.log.warning( + "Cannot deserialize conn_uri for connection '%s': upgrade to Airflow 3.2+ " + "or store the connection using individual fields (conn_type, host, login, " + "etc.) in Vault.", + conn_id, + ) + return None + def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | None: """ Retrieve a connection from Vault as a serialized string. From c948bf7a29489382a2bb89b6fcbed8e73a32a6d9 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Mon, 29 Jun 2026 09:59:48 -0500 Subject: [PATCH 8/9] Override get_conn_value() instead of get_connection() in VaultBackend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Return conn_uri verbatim or json.dumps(secret_dict) from get_conn_value() so the base-class get_connection() deserializes using the Connection class injected per execution context — fixing the mapper-init regression in workers without regressing server-side CLI commands. Tests assert on get_conn_value() string output directly, removing the _set_connection_class injection pattern. Add both conn_uri and field-based JWT auth tests to cover URI secrets. --- .../providers/hashicorp/secrets/vault.py | 38 ---- .../unit/hashicorp/secrets/test_vault.py | 180 ++++++++---------- 2 files changed, 77 insertions(+), 141 deletions(-) diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index 8d83e17bec87c..85fbb5dfad0e9 100644 --- a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py +++ b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py @@ -224,44 +224,6 @@ def _get_team_or_global_secret(self, base_path: str | None, team_name: str | Non return self._get_secret_with_base(path, key) - def get_connection(self, conn_id: str, team_name: str | None = None): - """ - Return connection object with a given ``conn_id``. - - Overrides the base to accept ``team_name`` on all supported Airflow versions and - to fall back to the compat SDK Connection when the framework-level class injection - (``deserialize_connection`` / ``_get_connection_class``) is not yet available. - - :param conn_id: connection id - :param team_name: Team name associated to the task trying to access the connection (if any) - :return: Connection object or None - """ - value = self.get_conn_value(conn_id=conn_id, team_name=team_name) - if value is None: - return None - # Prefer the base-class deserializer when available — it uses _get_connection_class() - # which the framework populates with the right Connection class per execution context. - if hasattr(self, "deserialize_connection"): - return self.deserialize_connection(conn_id=conn_id, value=value) - # Fallback for older Airflow that predates _get_connection_class / deserialize_connection. - from airflow.providers.common.compat.sdk import Connection - - value = value.strip() - if value.startswith("{"): - return Connection.from_json(value=value, conn_id=conn_id) - if hasattr(Connection, "from_uri"): - return Connection.from_uri(uri=value, conn_id=conn_id) - try: - return Connection(conn_id=conn_id, uri=value) # type: ignore[call-arg] - except (TypeError, ValueError): - self.log.warning( - "Cannot deserialize conn_uri for connection '%s': upgrade to Airflow 3.2+ " - "or store the connection using individual fields (conn_type, host, login, " - "etc.) in Vault.", - conn_id, - ) - return None - def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | None: """ Retrieve a connection from Vault as a serialized string. diff --git a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py index 4a2848587fc41..b6f1b8e3b1e0a 100644 --- a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py +++ b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py @@ -21,7 +21,6 @@ import pytest from hvac.exceptions import InvalidPath, VaultError -from airflow.providers.common.compat.sdk import Connection as SdkConnection from airflow.providers.hashicorp.secrets.vault import VaultBackend from tests_common.test_utils.config import conf_vars @@ -83,7 +82,9 @@ 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): + import json + mock_client = mock.MagicMock() mock_hvac.Client.return_value = mock_client mock_client.secrets.kv.v2.read_secret_version.return_value = { @@ -122,10 +123,12 @@ def test_get_connection(self, mock_hvac): } test_client = VaultBackend(**kwargs) - if hasattr(test_client, "_set_connection_class"): - test_client._set_connection_class(SdkConnection) - connection = test_client.get_connection(conn_id="test_postgres") - assert connection.get_uri() == "postgres://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"), @@ -183,9 +186,7 @@ def test_get_connection_value_multi_team( ) test_client = VaultBackend(**kwargs) - if hasattr(test_client, "_set_connection_class"): - test_client._set_connection_class(SdkConnection) - 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( @@ -197,7 +198,10 @@ def test_get_connection_value_multi_team( for path in exp_paths ] ) - assert connection.get_uri() == "postgres://airflow:airflow@host:5432/airflow?foo=bar&baz=taz" + import json + + 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): @@ -238,18 +242,19 @@ def test_get_connection_without_predefined_mount_point(self, mock_hvac): "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS", } + import json + test_client = VaultBackend(**kwargs) - if hasattr(test_client, "_set_connection_class"): - test_client._set_connection_class(SdkConnection) - connection = test_client.get_connection(conn_id="airflow/test_postgres") - assert connection.get_uri() == "postgres://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") @@ -547,7 +552,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" @@ -559,7 +564,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" @@ -571,10 +576,51 @@ 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_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 = { + "request_id": "94011e25-f8dc-ec29-221b-1f9c1d9ad2ae", + "lease_id": "", + "renewable": False, + "lease_duration": 0, + "data": { + "data": {"conn_uri": "postgresql://airflow:airflow@host:5432/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 == "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(self, mock_hvac): + def test_jwt_auth_type_field_based(self, mock_hvac): + import json + mock_client = mock.MagicMock() mock_hvac.Client.return_value = mock_client mock_client.secrets.kv.v2.read_secret_version.return_value = { @@ -613,10 +659,9 @@ def test_jwt_auth_type(self, mock_hvac): } test_client = VaultBackend(**kwargs) - if hasattr(test_client, "_set_connection_class"): - test_client._set_connection_class(SdkConnection) - 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 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" ) @@ -705,7 +750,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") @@ -736,10 +781,8 @@ def test_get_connection_with_empty_connections_path(self, mock_hvac): } backend = VaultBackend(**kwargs) - if hasattr(backend, "_set_connection_class"): - backend._set_connection_class(SdkConnection) - 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( @@ -749,7 +792,13 @@ 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" + import json + + 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): @@ -784,78 +833,3 @@ def test_config_path_none_value(self, mock_hvac): test_client = VaultBackend(**kwargs) assert test_client.get_config("test") is None mock_hvac.Client.assert_not_called() - - @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") - @pytest.mark.skipif( - not hasattr(VaultBackend, "_set_connection_class"), - reason="Connection class injection requires BaseSecretsBackend._set_connection_class (Airflow 3.2+)", - ) - def test_get_connection_returns_sdk_connection_not_sqlalchemy_model(self, mock_hvac): - """get_connection must return the SDK (Pydantic) Connection, not the SQLAlchemy ORM model. - - Using the SQLAlchemy model triggers lazy mapper initialisation for the entire Airflow model - registry, which fails in PythonVirtualenvOperator subprocesses where DagModel has not been - imported (raises sqlalchemy.exc.InvalidRequestError, silently swallowed in context.py). - """ - mock_client = mock.MagicMock() - mock_hvac.Client.return_value = mock_client - mock_client.secrets.kv.v2.read_secret_version.return_value = { - "data": { - "data": { - "conn_type": "postgres", - "host": "db-host", - "login": "user", - "password": "pass", - "port": "5432", - "schema": "mydb", - }, - "metadata": {"version": 1}, - } - } - - backend = VaultBackend( - connections_path="connections", - mount_point="airflow", - auth_type="token", - url="http://127.0.0.1:8200", - token="token", - ) - backend._set_connection_class(SdkConnection) - conn = backend.get_connection("trino_default") - - assert isinstance(conn, SdkConnection), ( - f"Expected SDK Connection, got {type(conn)}. " - "Returning the SQLAlchemy model triggers mapper init and breaks PythonVirtualenvOperator." - ) - assert conn.conn_id == "trino_default" - - @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") - @pytest.mark.skipif( - not hasattr(SdkConnection, "from_uri") or not hasattr(VaultBackend, "_set_connection_class"), - reason="conn_uri deserialization requires Connection.from_uri and _set_connection_class (Airflow 3.2+)", - ) - def test_get_connection_via_uri_returns_sdk_connection(self, mock_hvac): - mock_client = mock.MagicMock() - mock_hvac.Client.return_value = mock_client - mock_client.secrets.kv.v2.read_secret_version.return_value = { - "data": { - "data": {"conn_uri": "postgresql://user:pass@host:5432/db"}, - "metadata": {"version": 1}, - } - } - - backend = VaultBackend( - connections_path="connections", - mount_point="airflow", - auth_type="token", - url="http://127.0.0.1:8200", - token="token", - ) - backend._set_connection_class(SdkConnection) - conn = backend.get_connection("my_conn") - - assert isinstance(conn, SdkConnection), ( - f"Expected SDK Connection, got {type(conn)}. " - "Returning the SQLAlchemy model triggers mapper init and breaks PythonVirtualenvOperator." - ) - assert conn.conn_id == "my_conn" From 59b4285e77af5a295a67ee16f64fa99096498c93 Mon Sep 17 00:00:00 2001 From: seanmuth Date: Wed, 8 Jul 2026 14:22:51 -0500 Subject: [PATCH 9/9] Address review feedback on VaultBackend get_conn_value Hoist the repeated in-test json import to module level and qualify the get_conn_value docstring so the framework-injected Connection class behavior is scoped to Airflow 3.2+, where per-process injection exists (2.11/3.0/3.1 still import the ORM Connection directly). --- .../src/airflow/providers/hashicorp/secrets/vault.py | 11 +++++++---- .../tests/unit/hashicorp/secrets/test_vault.py | 11 +---------- 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index 85fbb5dfad0e9..8200c23c506ea 100644 --- a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py +++ b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py @@ -229,11 +229,14 @@ def get_conn_value(self, conn_id: str, team_name: str | None = None) -> str | No Retrieve a connection from Vault as a serialized string. Returns the ``conn_uri`` value verbatim when present, otherwise serializes - the secret dict to JSON. The base-class ``get_connection`` deserializes the - returned string using the Connection class that the framework injected for the - current execution context (ORM Connection on the server, SDK Connection in + 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. + 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. :param conn_id: connection id :param team_name: Team name associated to the task trying to access the connection (if any) diff --git a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py index b6f1b8e3b1e0a..53dc52a3f39e7 100644 --- a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py +++ b/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 @@ -83,8 +84,6 @@ def variable_result(self): @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_get_conn_value(self, mock_hvac): - import json - mock_client = mock.MagicMock() mock_hvac.Client.return_value = mock_client mock_client.secrets.kv.v2.read_secret_version.return_value = { @@ -198,8 +197,6 @@ def test_get_connection_value_multi_team( for path in exp_paths ] ) - import json - assert value is not None assert json.loads(value)["conn_type"] == "postgresql" @@ -242,8 +239,6 @@ def test_get_connection_without_predefined_mount_point(self, mock_hvac): "token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS", } - import json - test_client = VaultBackend(**kwargs) value = test_client.get_conn_value(conn_id="airflow/test_postgres") assert value is not None @@ -619,8 +614,6 @@ def test_jwt_auth_type_conn_uri(self, mock_hvac): @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_jwt_auth_type_field_based(self, mock_hvac): - import json - mock_client = mock.MagicMock() mock_hvac.Client.return_value = mock_client mock_client.secrets.kv.v2.read_secret_version.return_value = { @@ -792,8 +785,6 @@ def test_get_connection_with_empty_connections_path(self, mock_hvac): raise_on_deleted_version=True, ) - import json - assert value is not None parsed = json.loads(value) assert parsed["conn_type"] == "postgres"