[airflow] branch main updated: Invalidate Vault cached prop when not authenticated (#17387)

[ka...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kaxilnaik 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 16b47ce  Invalidate Vault cached prop when not authenticated (#17387)
16b47ce is described below

commit 16b47cecfb5cf88b0176a59589cbd77e0eaccfd3
Author: David Lum <da...@gmail.com>
AuthorDate: Wed Sep 1 09:53:04 2021 -0400

    Invalidate Vault cached prop when not authenticated (#17387)
---
 .../hashicorp/_internal_client/vault_client.py     | 19 ++++++++++++++-
 .../_internal_client/test_vault_client.py          | 27 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/airflow/providers/hashicorp/_internal_client/vault_client.py b/airflow/providers/hashicorp/_internal_client/vault_client.py
index d046c84..38c5c25 100644
--- a/airflow/providers/hashicorp/_internal_client/vault_client.py
+++ b/airflow/providers/hashicorp/_internal_client/vault_client.py
@@ -191,8 +191,25 @@ class _VaultClient(LoggingMixin):
         self.radius_secret = radius_secret
         self.radius_port = radius_port
 
+    @property
+    def client(self):
+        """
+        Authentication to Vault can expire. This wrapper function checks that
+        it is still authenticated to Vault, and invalidates the cache if this
+        is not the case.
+
+        :rtype: hvac.Client
+        :return: Vault Client
+
+        """
+        if not self._client.is_authenticated():
+            # Invalidate the cache:
+            # https://github.com/pydanny/cached-property#invalidating-the-cache
+            self.__dict__.pop('_client', None)
+        return self._client
+
     @cached_property
-    def client(self) -> hvac.Client:
+    def _client(self) -> hvac.Client:
         """
         Return an authenticated Hashicorp Vault client.
 
diff --git a/tests/providers/hashicorp/_internal_client/test_vault_client.py b/tests/providers/hashicorp/_internal_client/test_vault_client.py
index f59b860..f59872e 100644
--- a/tests/providers/hashicorp/_internal_client/test_vault_client.py
+++ b/tests/providers/hashicorp/_internal_client/test_vault_client.py
@@ -1003,3 +1003,30 @@ class TestVaultClient(TestCase):
         mock_client.secrets.kv.v1.create_or_update_secret.assert_called_once_with(
             mount_point='secret', secret_path='path', secret={'key': 'value'}, method="post"
         )
+
+    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
+    def test_cached_property_invalidates_on_auth_failure(self, mock_hvac):
+        mock_client = mock.MagicMock()
+        mock_hvac.Client.return_value = mock_client
+
+        vault_client = _VaultClient(
+            auth_type="radius",
+            radius_host="radhost",
+            radius_port=8110,
+            radius_secret="pass",
+            kv_engine_version=1,
+            url="http://localhost:8180",
+        )
+
+        # Assert that the original mock_client is returned
+        assert vault_client.client == mock_client
+
+        # Prove that the mock_client is cached by changing the return
+        # value, but still receive the original mock client
+        mock_client_2 = mock.MagicMock()
+        mock_hvac.Client.return_value = mock_client_2
+        assert vault_client.client == mock_client
+        mock_client.is_authenticated.return_value = False
+        # assert that when the client is not authenticated the cache
+        # is invalidated, therefore returning the second client
+        assert vault_client.client == mock_client_2