diff --git a/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py b/providers/hashicorp/src/airflow/providers/hashicorp/secrets/vault.py index 72ac30901b2e3..8200c23c506ea 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 @@ 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: """ diff --git a/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py b/providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py index 8b6af85a725da..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 @@ -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):