This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4112a4aa4ff Fix `VaultBackend.get_connection()` breaking
`PythonVirtualenvOperator` in Airflow 3 (#68305)
4112a4aa4ff is described below
commit 4112a4aa4fff06084df2bbd03055c8e069ea8653
Author: Sean Muth <[email protected]>
AuthorDate: Tue Jul 14 01:16:15 2026 -0500
Fix `VaultBackend.get_connection()` breaking `PythonVirtualenvOperator` in
Airflow 3 (#68305)
---
.../airflow/providers/hashicorp/secrets/vault.py | 33 +++----
.../tests/unit/hashicorp/secrets/test_vault.py | 101 +++++++++++++++++----
2 files changed, 100 insertions(+), 34 deletions(-)
diff --git
a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py
b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py
index 72ac30901b2..8200c23c506 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,32 +224,33 @@ class VaultBackend(BaseSecretsBackend, LoggingMixin):
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:
"""
diff --git a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py
b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py
index 8b6af85a725..53dc52a3f39 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
@@ -82,7 +83,7 @@ class TestVaultSecrets:
}
@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 @@ class TestVaultSecrets:
}
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 @@ class TestVaultSecrets:
)
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 @@ class TestVaultSecrets:
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 @@ class TestVaultSecrets:
}
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 @@ class TestVaultSecrets:
}
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 @@ class TestVaultSecrets:
}
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 @@ class TestVaultSecrets:
}
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 @@ class TestVaultSecrets:
}
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 @@ class TestVaultSecrets:
}
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 @@ class TestVaultSecrets:
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 @@ class TestVaultSecrets:
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 @@ class TestVaultSecrets:
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):
