From 4f22ad017229cb3b956007ca8acc1ac45f997d19 Mon Sep 17 00:00:00 2001
From: Pablo Martin <pablo.martin@lolamarket.com>
Date: Mon, 16 Jan 2023 13:29:44 +0100
Subject: [PATCH] Refactor the fetching of env data with a specific function
 and a decorator

---
 lolafect/lolaconfig.py | 48 ++++++++++++++++++++++++------------------
 1 file changed, 28 insertions(+), 20 deletions(-)

diff --git a/lolafect/lolaconfig.py b/lolafect/lolaconfig.py
index 786855a..a8265eb 100644
--- a/lolafect/lolaconfig.py
+++ b/lolafect/lolaconfig.py
@@ -14,6 +14,14 @@ from lolafect.defaults import (
 from lolafect.utils import S3FileReader
 
 
+def _needs_env_data(method):
+    def wrapper(self, *args, **kwargs):
+        self.fetch_env_data(kwargs.get("s3_reader", None))
+        return method(self, *args, **kwargs)
+
+    return wrapper
+
+
 class LolaConfig:
     """
     A global-ish container for configurations required in pretty much all flows.
@@ -68,6 +76,7 @@ class LolaConfig:
             DEFAULT_KUBERNETES_IMAGE if kubernetes_image is None else kubernetes_image
         )
 
+        self.ENV_DATA = None
         self.TRINO_CREDENTIALS = None
         self.SSH_TUNNEL_CREDENTIALS = None
 
@@ -88,6 +97,7 @@ class LolaConfig:
             bucket=self.S3_BUCKET_NAME, key=self.SLACK_WEBHOOKS_FILE
         )
 
+    @_needs_env_data
     def fetch_trino_credentials(self, s3_reader=None) -> None:
         """
         Read the env file from S3 and store the trino credentials in memory.
@@ -96,20 +106,14 @@ class LolaConfig:
         :return: None
         """
 
-        if s3_reader is None:
-            s3_reader = self._s3_reader
-
-        env_data = s3_reader.read_json_from_s3_file(
-            bucket=self.S3_BUCKET_NAME, key=self.ENV_FILE_PATH
-        )
-
         self.TRINO_CREDENTIALS = {
-            "host": env_data["trino_host"],
-            "user": env_data["trino_user"],
-            "password": env_data["trino_pass"],
-            "port": env_data["trino_port"],
+            "host": self.ENV_DATA["trino_host"],
+            "user": self.ENV_DATA["trino_user"],
+            "password": self.ENV_DATA["trino_pass"],
+            "port": self.ENV_DATA["trino_port"],
         }
 
+    @_needs_env_data
     def fetch_ssh_tunnel_credentials(self, s3_reader=None) -> None:
         """
         Read the env file from S3 and store the SSH tunnel credentials.
@@ -118,20 +122,24 @@ class LolaConfig:
         :return: None
         """
 
+        self.SSH_TUNNEL_CREDENTIALS = {
+            "path_to_ssh_pkey": self.ENV_DATA["pt_ssh_pkey_path"],
+            "ssh_pkey_password": self.ENV_DATA["pt_ssh_pkey_passphrase"],
+            "ssh_username": self.ENV_DATA["pt_ssh_username"],
+            "ssh_port": self.ENV_DATA["pt_ssh_jumphost_port"],
+            "ssh_jumphost": self.ENV_DATA["pt_ssh_jumphost"],
+        }
+
+    def fetch_env_data(self, s3_reader=None):
+        if self.ENV_DATA is not None:
+            return
+
         if s3_reader is None:
             s3_reader = self._s3_reader
-
         env_data = s3_reader.read_json_from_s3_file(
             bucket=self.S3_BUCKET_NAME, key=self.ENV_FILE_PATH
         )
-
-        self.SSH_TUNNEL_CREDENTIALS = {
-            "path_to_ssh_pkey": env_data["pt_ssh_pkey_path"],
-            "ssh_pkey_password": env_data["pt_ssh_pkey_passphrase"],
-            "ssh_username": env_data["pt_ssh_username"],
-            "ssh_port": env_data["pt_ssh_jumphost_port"],
-            "ssh_jumphost": env_data["pt_ssh_jumphost"],
-        }
+        self.ENV_DATA = env_data
 
 
 def build_lolaconfig(