From cd9ba6137ffbddd57519d98f3bda78b16181b0ce Mon Sep 17 00:00:00 2001
From: Samson Sebastian <ssebastian@clearstreet.io>
Date: Mon, 8 Jan 2024 13:26:40 -0500
Subject: [PATCH] Removes private methods for init and get. Retries only on
 vault_client calls. Timeout set to 60 seconds, retries to 5.

---
 gestalt/vault.py | 66 ++++++++++++++++++------------------------------
 1 file changed, 24 insertions(+), 42 deletions(-)

diff --git a/gestalt/vault.py b/gestalt/vault.py
index 331ba12..e1a914b 100644
--- a/gestalt/vault.py
+++ b/gestalt/vault.py
@@ -24,30 +24,8 @@ def __init__(
         token: Optional[str] = os.environ.get("VAULT_TOKEN"),
         verify: Optional[bool] = True,
         scheme: str = "ref+vault://",
-        delay: int = 2,
+        delay: int = 60,
         tries: int = 5,
-    ) -> None:
-        self.delay = delay
-        self.tries = tries
-
-        retry_call(
-            f=Vault.__do_init,
-            fargs=[self, cert, role, jwt, url, token, verify, scheme],
-            exceptions=(RuntimeError, Timeout),
-            delay=self.delay,
-            tries=self.tries,
-        )
-
-    # __init__ impl for retry_call
-    def __do_init(
-            self,
-            cert: Optional[Tuple[str, str]],
-            role: Optional[str],
-            jwt: Optional[str],
-            url: Optional[str],
-            token: Optional[str],
-            verify: Optional[bool],
-            scheme: str,
     ) -> None:
         """Initialized vault client and authenticates vault
         Args:
@@ -70,8 +48,16 @@ def __do_init(
         self._secret_values: Dict[str, Union[str, int, float, bool,
                                              List[Any]]] = dict()
 
+        self.delay = delay
+        self.tries = tries
+
         try:
-            self.vault_client.is_authenticated()
+            retry_call(
+                self.vault_client.is_authenticated,
+                exceptions=(RuntimeError, Timeout),
+                delay=self.delay,
+                tries=self.tries,
+            )
         except requests.exceptions.MissingSchema:
             raise RuntimeError(
                 "Gestalt Error: Unable to connect to vault with the given configuration"
@@ -81,7 +67,13 @@ def __do_init(
             try:
                 hvac.api.auth_methods.Kubernetes(
                     self.vault_client.adapter).login(role=role, jwt=jwt)
-                token = self.vault_client.auth.token.lookup_self()
+                token = retry_call(
+                    self.vault_client.auth.token.lookup_self,
+                    exceptions=(RuntimeError, Timeout),
+                    delay=self.delay,
+                    tries=self.tries,
+                )
+
                 if token is not None:
                     kubes_token = (
                         "kubernetes",
@@ -121,22 +113,6 @@ def get(
         path: str,
         filter: str,
         sep: Optional[str] = "."
-    ) -> Union[str, int, float, bool, List[Any]]:
-        return retry_call(
-            f=Vault.__do_get,
-            fargs=[self, key, path, filter, sep],
-            exceptions=(RuntimeError, Timeout),
-            delay=self.delay,
-            tries=self.tries,
-        )
-
-    # get impl for retry_call
-    def __do_get(
-        self,
-        key: str,
-        path: str,
-        filter: str,
-        sep: Optional[str] = "."
     ) -> Union[str, int, float, bool, List[Any]]:
         """Gets secret from vault
         Args:
@@ -157,7 +133,13 @@ def __do_get(
             return self._secret_values[key]
 
         try:
-            response = self.vault_client.read(path)
+            response = retry_call(
+                self.vault_client.read,
+                fargs=[path],
+                exceptions=(RuntimeError, Timeout),
+                delay=self.delay,
+                tries=self.tries,
+            )
             if response is None:
                 raise RuntimeError("Gestalt Error: No secrets found")
             if response["lease_id"]: