Lazily connect to vault (#43)

* lazily connect

* lazily connect

* lazily connect

* add to changelog

* edit tests

* tests

* debugging

* debugging

* linting

* no diff

* testing

* testing

* testing

* testing

* testing

* formatting

* testing

* add unit test

* tests

* tests

* tests

* tests

* add property

* add property

* formatting

* formatting

CHANGELOG.md

@@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 
 
+## [3.4.0] - 2024-03-04
+
+### Added
+- Lazily authenticates to vault upon first vault request, instead of authentication on object creation.
+
 ## [3.3.6] - 2023-06-11
 
 ### Added

gestalt/vault.py

@@ -41,13 +41,31 @@ def __init__(
         self.dynamic_token_queue: Queue[Tuple[str, str, str]] = Queue()
         self.kubes_token_queue: Queue[Tuple[str, str, str]] = Queue()
 
-        self.vault_client = hvac.Client(url=url, token=token, cert=cert, verify=verify)
+        self._vault_client: Optional[hvac.Client] = None
         self._secret_expiry_times: Dict[str, datetime] = dict()
-        self._secret_values: Dict[str, Union[str, int, float, bool, List[Any]]] = dict()
+        self._secret_values: Dict[str, Union[str, int, float, bool,
+                                             List[Any]]] = dict()
+        self._is_connected: bool = False
+        self._role: Optional[str] = role
+        self._jwt: Optional[str] = jwt
+        self._url: Optional[str] = url
+        self._token: Optional[str] = token
+        self._cert: Optional[Tuple[str, str]] = cert
+        self._verify: Optional[bool] = verify
 
         self.delay = delay
         self.tries = tries
 
+    @property
+    def vault_client(self) -> hvac.Client:
+        if self._vault_client is None:
+            self._vault_client = hvac.Client(url=self._url,
+                                             token=self._token,
+                                             cert=self._cert,
+                                             verify=self._verify)
+        return self._vault_client
+
+    def connect(self) -> None:
         try:
             retry_call(
                 self.vault_client.is_authenticated,
@@ -60,11 +78,11 @@ def __init__(
                 "Gestalt Error: Unable to connect to vault with the given configuration"
             )
 
-        if role and jwt:
+        if self._role and self._jwt:
             try:
-                hvac.api.auth_methods.Kubernetes(self.vault_client.adapter).login(
-                    role=role, jwt=jwt
-                )
+                hvac.api.auth_methods.Kubernetes(
+                    self.vault_client.adapter).login(role=self._role,
+                                                     jwt=self._jwt)
                 token = retry_call(
                     self.vault_client.auth.token.lookup_self,
                     exceptions=(RuntimeError, Timeout),
@@ -75,30 +93,30 @@ def __init__(
                 if token is not None:
                     kubes_token = (
                         "kubernetes",
-                        token["data"]["id"],  # type: ignore
+                        token["data"]["id"],
                         token["data"]["ttl"],
-                    )  # type: ignore
+                    )
                     self.kubes_token_queue.put(kubes_token)
             except hvac.exceptions.InvalidPath:
                 raise RuntimeError(
-                    "Gestalt Error: Kubernetes auth couldn't be performed"
-                )
+                    "Gestalt Error: Kubernetes auth couldn't be performed")
             except requests.exceptions.ConnectionError:
                 raise RuntimeError("Gestalt Error: Couldn't connect to Vault")
 
             dynamic_ttl_renew = Thread(
                 name="dynamic-token-renew",
                 target=self.worker,
                 daemon=True,
-                args=(self.dynamic_token_queue,),
+                args=(self.dynamic_token_queue, ),
             )  # noqa: F841
             kubernetes_ttl_renew = Thread(
                 name="kubes-token-renew",
                 target=self.worker,
                 daemon=True,
-                args=(self.kubes_token_queue,),
+                args=(self.kubes_token_queue, ),
             )
             kubernetes_ttl_renew.start()
+        self._is_connected = True
 
     def stop(self) -> None:
         self._run_worker = False
@@ -107,7 +125,11 @@ def __del__(self) -> None:
         self.stop()
 
     def get(
-        self, key: str, path: str, filter: str, sep: Optional[str] = "."
+        self,
+        key: str,
+        path: str,
+        filter: str,
+        sep: Optional[str] = "."
     ) -> Union[str, int, float, bool, List[Any]]:
         """Gets secret from vault
         Args:
@@ -118,12 +140,15 @@ def get(
         Returns:
             secret (str): secret
         """
+        if not self._is_connected:
+            self.connect()
         # if the key has been read before and is not a TTL secret
         if key in self._secret_values and key not in self._secret_expiry_times:
             return self._secret_values[key]
 
         # if the secret can expire but hasn't expired yet
-        if key in self._secret_expiry_times and not self._is_secret_expired(key):
+        if key in self._secret_expiry_times and not self._is_secret_expired(
+                key):
             return self._secret_values[key]
 
         try:
@@ -146,10 +171,10 @@ def get(
             requested_data = response["data"].get("data", response["data"])
         except hvac.exceptions.InvalidPath:
             raise RuntimeError(
-                "Gestalt Error: The secret path or mount is set incorrectly"
-            )
+                "Gestalt Error: The secret path or mount is set incorrectly")
         except requests.exceptions.ConnectionError:
-            raise RuntimeError("Gestalt Error: Gestalt couldn't connect to Vault")
+            raise RuntimeError(
+                "Gestalt Error: Gestalt couldn't connect to Vault")
         except Exception as err:
             raise RuntimeError(f"Gestalt Error: {err}")
         if filter is None:
@@ -175,13 +200,12 @@ def _is_secret_expired(self, key: str) -> bool:
         is_expired = now >= secret_expires_dt
         return is_expired
 
-    def _set_secrets_ttl(self, requested_data: Dict[str, Any], key: str) -> None:
-        last_vault_rotation_str = requested_data["last_vault_rotation"].split(".")[
-            0
-        ]  # to the nearest second
-        last_vault_rotation_dt = datetime.strptime(
-            last_vault_rotation_str, "%Y-%m-%dT%H:%M:%S"
-        )
+    def _set_secrets_ttl(self, requested_data: Dict[str, Any],
+                         key: str) -> None:
+        last_vault_rotation_str = requested_data["last_vault_rotation"].split(
+            ".")[0]  # to the nearest second
+        last_vault_rotation_dt = datetime.strptime(last_vault_rotation_str,
+                                                   "%Y-%m-%dT%H:%M:%S")
         ttl = requested_data["ttl"]
         secret_expires_dt = last_vault_rotation_dt + timedelta(seconds=ttl)
         self._secret_expiry_times[key] = secret_expires_dt
@@ -190,11 +214,11 @@ def worker(self, token_queue: Queue) -> None:  # type: ignore
         """
         Worker function to renew lease on expiry
         """
-
         try:
             while self._run_worker:
                 if not token_queue.empty():
-                    token_type, token_id, token_duration = token = token_queue.get()
+                    token_type, token_id, token_duration = token = token_queue.get(
+                    )
                     if token_type == "kubernetes":
                         self.vault_client.auth.token.renew(token_id)
                         print("kubernetes token for the app has been renewed")
@@ -206,10 +230,10 @@ def worker(self, token_queue: Queue) -> None:  # type: ignore
                     sleep((token_duration / 3) * 2)
         except hvac.exceptions.InvalidPath:
             raise RuntimeError(
-                "Gestalt Error: The lease path or mount is set incorrectly"
-            )
+                "Gestalt Error: The lease path or mount is set incorrectly")
         except requests.exceptions.ConnectionError:
-            raise RuntimeError("Gestalt Error: Gestalt couldn't connect to Vault")
+            raise RuntimeError(
+                "Gestalt Error: Gestalt couldn't connect to Vault")
         except Exception as err:
             raise RuntimeError(f"Gestalt Error: {err}")
 

setup.py

@@ -2,35 +2,37 @@
 
 
 def readme():
-    with open('README.md') as f:
+    with open("README.md") as f:
         return f.read()
 
 
 with open("requirements.txt") as reqs_file:
     reqs = filter(lambda x: not x.startswith("-"), reqs_file.readlines())
     reqs_list = list(map(lambda x: x.rstrip(), reqs))
 
-setup(name='gestalt-cfg',
-      version='3.3.6',
-      description='A sensible configuration library for Python',
-      long_description=readme(),
-      long_description_content_type="text/markdown",
-      url='https://github.com/clear-street/gestalt',
-      author='Clear Street',
-      author_email='[email protected]',
-      license='MIT',
-      packages=find_packages(),
-      python_requires='>=3.6',
-      install_requires=reqs_list,
-      classifiers=[
-          'Development Status :: 5 - Production/Stable',
-          'Intended Audience :: Developers',
-          'License :: OSI Approved :: MIT License',
-          'Topic :: Software Development :: Libraries',
-          'Programming Language :: Python :: 3',
-          'Programming Language :: Python :: 3.6',
-          'Programming Language :: Python :: 3.7',
-          'Programming Language :: Python :: 3.8',
-          'Programming Language :: Python :: 3 :: Only',
-          'Operating System :: OS Independent',
-      ])
+setup(
+    name="gestalt-cfg",
+    version="3.4.0",
+    description="A sensible configuration library for Python",
+    long_description=readme(),
+    long_description_content_type="text/markdown",
+    url="https://github.com/clear-street/gestalt",
+    author="Clear Street",
+    author_email="[email protected]",
+    license="MIT",
+    packages=find_packages(),
+    python_requires=">=3.6",
+    install_requires=reqs_list,
+    classifiers=[
+        "Development Status :: 5 - Production/Stable",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: MIT License",
+        "Topic :: Software Development :: Libraries",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.6",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3 :: Only",
+        "Operating System :: OS Independent",
+    ],
+)

tests/test_gestalt.py

@@ -557,6 +557,15 @@ def test_set_vault_key(nested_setup):
     assert secret == "ref+vault://secret/data/testnested#.slack.token"
 
 
+def test_vault_lazy_connect(mock_vault_workers, mock_vault_k8s_auth):
+    with patch("gestalt.vault.hvac.Client") as mock_client:
+        v = Vault(role="test-role", jwt="test-jwt")
+        assert not v._is_connected
+        v.get("foo", "foo", ".foo")
+        assert v._is_connected
+        mock_client().auth.token.lookup_self.assert_called()
+
+
 def test_vault_worker_dynamic(mock_vault_workers, mock_vault_k8s_auth):
     mock_dynamic_renew, mock_k8s_renew = mock_vault_workers
 
@@ -571,6 +580,7 @@ def except_once(self, **kwargs):
                autospec=True) as mock_sleep:
         with patch("gestalt.vault.hvac.Client") as mock_client:
             v = Vault(role="test-role", jwt="test-jwt")
+            v.connect()
 
             mock_k8s_renew.start.assert_called()
 
@@ -602,6 +612,7 @@ def except_once(self, **kwargs):
                autospec=True) as mock_sleep:
         with patch("gestalt.vault.hvac.Client") as mock_client:
             v = Vault(role="test-role", jwt="test-jwt")
+            v.connect()
 
             mock_k8s_renew.start.assert_called()