diff --git a/airflow/providers/hashicorp/_internal_client/vault_client.py b/airflow/providers/hashicorp/_internal_client/vault_client.py index aea8bfb01d76e..357c80d70e5db 100644 --- a/airflow/providers/hashicorp/_internal_client/vault_client.py +++ b/airflow/providers/hashicorp/_internal_client/vault_client.py @@ -16,6 +16,13 @@ # under the License. from __future__ import annotations +import sys + +if sys.version_info < (3, 8): + from importlib_metadata import version +else: + from importlib.metadata import version + from functools import cached_property import hvac @@ -30,7 +37,6 @@ DEFAULT_KUBERNETES_JWT_PATH = "/var/run/secrets/kubernetes.io/serviceaccount/token" DEFAULT_KV_ENGINE_VERSION = 2 - VALID_KV_VERSIONS: list[int] = [1, 2] VALID_AUTH_TYPES: list[str] = [ "approle", @@ -365,6 +371,7 @@ def get_secret(self, secret_path: str, secret_version: int | None = None) -> dic :return: secret stored in the vault as a dictionary """ mount_point = None + hvac_version = version("hvac") try: mount_point, secret_path = self._parse_secret_path(secret_path) if self.kv_engine_version == 1: @@ -372,9 +379,17 @@ def get_secret(self, secret_path: str, secret_version: int | None = None) -> dic raise VaultError("Secret version can only be used with version 2 of the KV engine") response = self.client.secrets.kv.v1.read_secret(path=secret_path, mount_point=mount_point) else: - response = self.client.secrets.kv.v2.read_secret_version( - path=secret_path, mount_point=mount_point, version=secret_version - ) + if hvac_version >= "1.1.0": + response = self.client.secrets.kv.v2.read_secret_version( + path=secret_path, + mount_point=mount_point, + version=secret_version, + raise_on_deleted_version=True, + ) + else: + response = self.client.secrets.kv.v2.read_secret_version( + path=secret_path, mount_point=mount_point, version=secret_version + ) except InvalidPath: self.log.debug("Secret not found %s with mount point %s", secret_path, mount_point) return None @@ -419,11 +434,20 @@ def get_secret_including_metadata( if self.kv_engine_version == 1: raise VaultError("Metadata might only be used with version 2 of the KV engine.") mount_point = None + hvac_version = version("hvac") try: mount_point, secret_path = self._parse_secret_path(secret_path) - return self.client.secrets.kv.v2.read_secret_version( - path=secret_path, mount_point=mount_point, version=secret_version - ) + if hvac_version >= "1.1.0": + return self.client.secrets.kv.v2.read_secret_version( + path=secret_path, + mount_point=mount_point, + version=secret_version, + raise_on_deleted_version=True, + ) + else: + return self.client.secrets.kv.v2.read_secret_version( + path=secret_path, mount_point=mount_point, version=secret_version + ) except InvalidPath: self.log.debug( "Secret not found %s with mount point %s and version %s", diff --git a/tests/providers/hashicorp/_internal_client/test_vault_client.py b/tests/providers/hashicorp/_internal_client/test_vault_client.py index bb9a53ceb5327..cd6e1bdb695e2 100644 --- a/tests/providers/hashicorp/_internal_client/test_vault_client.py +++ b/tests/providers/hashicorp/_internal_client/test_vault_client.py @@ -16,6 +16,13 @@ # under the License. from __future__ import annotations +import sys + +if sys.version_info < (3, 8): + from importlib_metadata import version +else: + from importlib.metadata import version + from unittest import mock from unittest.mock import mock_open, patch @@ -640,9 +647,15 @@ def test_get_non_existing_key_v2(self, mock_hvac): ) secret = vault_client.get_secret(secret_path="missing") assert secret is None - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="secret", path="missing", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None + ) @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_get_non_existing_key_v2_different_auth(self, mock_hvac): @@ -660,9 +673,15 @@ def test_get_non_existing_key_v2_different_auth(self, mock_hvac): secret = vault_client.get_secret(secret_path="missing") assert secret is None assert "secret" == vault_client.mount_point - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="secret", path="missing", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None + ) @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_get_non_existing_key_v1(self, mock_hvac): @@ -715,9 +734,15 @@ def test_get_existing_key_v2(self, mock_hvac): ) secret = vault_client.get_secret(secret_path="path/to/secret") assert {"secret_key": "secret_value"} == secret - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="secret", path="path/to/secret", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="path/to/secret", version=None, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="path/to/secret", version=None + ) @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_get_existing_key_v2_without_preconfigured_mount_point(self, mock_hvac): @@ -753,9 +778,15 @@ def test_get_existing_key_v2_without_preconfigured_mount_point(self, mock_hvac): ) secret = vault_client.get_secret(secret_path="mount_point/path/to/secret") assert {"secret_key": "secret_value"} == secret - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="mount_point", path="path/to/secret", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="mount_point", path="path/to/secret", version=None, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="mount_point", path="path/to/secret", version=None + ) @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_get_existing_key_v2_version(self, mock_hvac): @@ -790,9 +821,15 @@ def test_get_existing_key_v2_version(self, mock_hvac): ) secret = vault_client.get_secret(secret_path="missing", secret_version=1) assert {"secret_key": "secret_value"} == secret - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="secret", path="missing", version=1 - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=1, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=1 + ) @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_get_existing_key_v1(self, mock_hvac): @@ -1014,9 +1051,15 @@ def test_get_secret_including_metadata_v2(self, mock_hvac): "warnings": None, "auth": None, } == metadata - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="secret", path="missing", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None + ) @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") def test_get_secret_including_metadata_v1(self, mock_hvac): diff --git a/tests/providers/hashicorp/hooks/test_vault.py b/tests/providers/hashicorp/hooks/test_vault.py index 4bd3e90e56146..3afa477670877 100644 --- a/tests/providers/hashicorp/hooks/test_vault.py +++ b/tests/providers/hashicorp/hooks/test_vault.py @@ -16,6 +16,13 @@ # under the License. from __future__ import annotations +import sys + +if sys.version_info < (3, 8): + from importlib_metadata import version +else: + from importlib.metadata import version + from unittest import mock from unittest.mock import PropertyMock, mock_open, patch @@ -1001,9 +1008,15 @@ def test_get_existing_key_v2(self, mock_hvac, mock_get_connection): test_hook = VaultHook(**kwargs) secret = test_hook.get_secret(secret_path="missing") assert {"secret_key": "secret_value"} == secret - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="secret", path="missing", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None + ) @mock.patch("airflow.providers.hashicorp.hooks.vault.VaultHook.get_connection") @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") @@ -1040,9 +1053,15 @@ def test_get_existing_key_v2_version(self, mock_hvac, mock_get_connection): test_hook = VaultHook(**kwargs) secret = test_hook.get_secret(secret_path="missing", secret_version=1) assert {"secret_key": "secret_value"} == secret - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="secret", path="missing", version=1 - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=1, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=1 + ) @mock.patch("airflow.providers.hashicorp.hooks.vault.VaultHook.get_connection") @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") @@ -1185,9 +1204,15 @@ def test_get_secret_including_metadata_v2(self, mock_hvac, mock_get_connection): "warnings": None, "auth": None, } == metadata - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="secret", path="missing", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="secret", path="missing", version=None + ) @mock.patch("airflow.providers.hashicorp.hooks.vault.VaultHook.get_connection") @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") diff --git a/tests/providers/hashicorp/secrets/test_vault.py b/tests/providers/hashicorp/secrets/test_vault.py index 4897a73c22334..58f04a2aa04d6 100644 --- a/tests/providers/hashicorp/secrets/test_vault.py +++ b/tests/providers/hashicorp/secrets/test_vault.py @@ -16,6 +16,13 @@ # under the License. from __future__ import annotations +import sys + +if sys.version_info < (3, 8): + from importlib_metadata import version +else: + from importlib.metadata import version + from unittest import mock import pytest @@ -301,9 +308,18 @@ def test_get_conn_uri_non_existent_key(self, mock_hvac): test_client = VaultBackend(**kwargs) assert test_client.get_conn_uri(conn_id="test_mysql") is None - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="airflow", path="connections/test_mysql", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="airflow", + path="connections/test_mysql", + version=None, + raise_on_deleted_version=True, + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="airflow", path="connections/test_mysql", version=None + ) assert test_client.get_connection(conn_id="test_mysql") is None @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac") @@ -453,9 +469,15 @@ def test_get_variable_value_non_existent_key(self, mock_hvac): test_client = VaultBackend(**kwargs) assert test_client.get_variable("hello") is None - mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( - mount_point="airflow", path="variables/hello", version=None - ) + hvac_version = version("hvac") + if hvac_version >= "1.1.0": + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="airflow", path="variables/hello", version=None, raise_on_deleted_version=True + ) + else: + mock_client.secrets.kv.v2.read_secret_version.assert_called_once_with( + mount_point="airflow", path="variables/hello", version=None + ) assert test_client.get_variable("hello") is None @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")