diff --git a/lolafect/lolaconfig.py b/lolafect/lolaconfig.py index 1a71a92..9ac66c0 100644 --- a/lolafect/lolaconfig.py +++ b/lolafect/lolaconfig.py @@ -9,7 +9,7 @@ from lolafect.defaults import ( DEFAULT_KUBERNETES_IMAGE, DEFAULT_KUBERNETES_LABELS, DEFAULT_FLOWS_PATH_IN_BUCKET, - DEFAULT_ENV_FILE_PATH + DEFAULT_ENV_FILE_PATH, ) from lolafect.utils import S3FileReader @@ -68,6 +68,8 @@ class LolaConfig: DEFAULT_KUBERNETES_IMAGE if kubernetes_image is None else kubernetes_image ) + self.TRINO_CREDENTIALS = None + self._s3_reader = S3FileReader(s3_client=boto3.client("s3")) def fetch_slack_webhooks(self, s3_reader=None) -> None: @@ -85,6 +87,28 @@ class LolaConfig: bucket=self.S3_BUCKET_NAME, key=self.SLACK_WEBHOOKS_FILE ) + def fetch_trino_credentials(self, s3_reader=None) -> None: + """ + Read the env file from S3 and store the trino credentials in memory. + + :param s3_reader: a client to fetch files from S3. + :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"], + } + def build_lolaconfig( flow_name: str, diff --git a/tests/test_lolaconfig.py b/tests/test_lolaconfig.py index 742de77..5d7189a 100644 --- a/tests/test_lolaconfig.py +++ b/tests/test_lolaconfig.py @@ -4,12 +4,10 @@ from lolafect.lolaconfig import LolaConfig def test_lolaconfig_instantiates_with_defaults_properly(): - lolaconfig = LolaConfig(flow_name="some-flow") def test_lolaconfig_instantiates_without_defaults_proplery(): - lolaconfig = LolaConfig( flow_name="some-flow", env_s3_bucket="bucket", @@ -20,7 +18,6 @@ def test_lolaconfig_instantiates_without_defaults_proplery(): def test_lolaconfig_fetches_webhooks_properly(): - lolaconfig = LolaConfig(flow_name="some-flow") fake_s3_reader = SimpleNamespace() @@ -33,3 +30,23 @@ def test_lolaconfig_fetches_webhooks_properly(): lolaconfig.fetch_slack_webhooks(s3_reader=fake_s3_reader) assert type(lolaconfig.SLACK_WEBHOOKS) is dict + + +def test_lolaconfig_fetches_trino_creds_properly(): + lolaconfig = LolaConfig(flow_name="some-flow") + + fake_s3_reader = SimpleNamespace() + + def mock_read_json_from_s3_file(bucket, key): + return { + "trino_host": "some_host", + "trino_user": "some_user", + "trino_pass": "some_password", + "trino_port": "some_port", + } + + fake_s3_reader.read_json_from_s3_file = mock_read_json_from_s3_file + + lolaconfig.fetch_trino_credentials(s3_reader=fake_s3_reader) + + assert type(lolaconfig.TRINO_CREDENTIALS) is dict