Refactor the fetching of env data with a specific function and a decorator

This commit is contained in:
Pablo Martin 2023-01-16 13:29:44 +01:00
parent 592dd6bb69
commit 4f22ad0172

View file

@ -14,6 +14,14 @@ from lolafect.defaults import (
from lolafect.utils import S3FileReader 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: class LolaConfig:
""" """
A global-ish container for configurations required in pretty much all flows. 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 DEFAULT_KUBERNETES_IMAGE if kubernetes_image is None else kubernetes_image
) )
self.ENV_DATA = None
self.TRINO_CREDENTIALS = None self.TRINO_CREDENTIALS = None
self.SSH_TUNNEL_CREDENTIALS = None self.SSH_TUNNEL_CREDENTIALS = None
@ -88,6 +97,7 @@ class LolaConfig:
bucket=self.S3_BUCKET_NAME, key=self.SLACK_WEBHOOKS_FILE bucket=self.S3_BUCKET_NAME, key=self.SLACK_WEBHOOKS_FILE
) )
@_needs_env_data
def fetch_trino_credentials(self, s3_reader=None) -> None: def fetch_trino_credentials(self, s3_reader=None) -> None:
""" """
Read the env file from S3 and store the trino credentials in memory. Read the env file from S3 and store the trino credentials in memory.
@ -96,20 +106,14 @@ class LolaConfig:
:return: None :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 = { self.TRINO_CREDENTIALS = {
"host": env_data["trino_host"], "host": self.ENV_DATA["trino_host"],
"user": env_data["trino_user"], "user": self.ENV_DATA["trino_user"],
"password": env_data["trino_pass"], "password": self.ENV_DATA["trino_pass"],
"port": env_data["trino_port"], "port": self.ENV_DATA["trino_port"],
} }
@_needs_env_data
def fetch_ssh_tunnel_credentials(self, s3_reader=None) -> None: def fetch_ssh_tunnel_credentials(self, s3_reader=None) -> None:
""" """
Read the env file from S3 and store the SSH tunnel credentials. Read the env file from S3 and store the SSH tunnel credentials.
@ -118,20 +122,24 @@ class LolaConfig:
:return: None :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: if s3_reader is None:
s3_reader = self._s3_reader s3_reader = self._s3_reader
env_data = s3_reader.read_json_from_s3_file( env_data = s3_reader.read_json_from_s3_file(
bucket=self.S3_BUCKET_NAME, key=self.ENV_FILE_PATH bucket=self.S3_BUCKET_NAME, key=self.ENV_FILE_PATH
) )
self.ENV_DATA = env_data
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"],
}
def build_lolaconfig( def build_lolaconfig(