diff --git a/README.md b/README.md index a2e40b0..5af6c0a 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,13 @@ # Lolafect -Lolafect is a collection of Python bits that help us build our Prefect flows. \ No newline at end of file +Lolafect is a collection of Python bits that help us build our Prefect flows. + + +## How to test + +IDE-agnostic: +1. Set up a virtual environment which contains both `lolafect` and the dependencies listed in `requirements-dev.txt`. +2. Run: `pytests tests` + +In Pycharm: if you configure `pytest` as the project test runner, Pycharm will most probably autodetect the test +folder and allow you to run the test suite within the IDE. \ No newline at end of file diff --git a/lolafect/defaults.py b/lolafect/defaults.py new file mode 100644 index 0000000..e8397bc --- /dev/null +++ b/lolafect/defaults.py @@ -0,0 +1,5 @@ +DEFAULT_ENV_S3_BUCKET="pdo-prefect-flows" +DEFAULT_PATH_TO_SLACK_WEBHOOKS_FILE = "env/slack_webhooks.json" +DEFAULT_KUBERNETES_IMAGE = "373245262072.dkr.ecr.eu-central-1.amazonaws.com/pdo-data-prefect:production" +DEFAULT_KUBERNETES_LABELS = ["k8s"] +DEFAULT_FLOWS_PATH_IN_BUCKET = "flows/" diff --git a/lolafect/lolaconfig.py b/lolafect/lolaconfig.py new file mode 100644 index 0000000..e7eefda --- /dev/null +++ b/lolafect/lolaconfig.py @@ -0,0 +1,109 @@ +from typing import List + +from prefect.storage import S3 +import boto3 + +from lolafect.defaults import ( + DEFAULT_ENV_S3_BUCKET, + DEFAULT_PATH_TO_SLACK_WEBHOOKS_FILE, + DEFAULT_KUBERNETES_IMAGE, + DEFAULT_KUBERNETES_LABELS, + DEFAULT_FLOWS_PATH_IN_BUCKET, +) +from lolafect.utils import S3FileReader + + +class LolaConfig: + """ + A global-ish container for configurations required in pretty much all flows. + """ + + def __init__( + self, + flow_name: str, + env_s3_bucket: str = None, + kubernetes_labels: List = None, + kubernetes_image: str = None, + slack_webhooks_file: str = None, + ): + """ + Init and set defaults where no value was passed. + + :param flow_name: the name of the flow. + :param env_s3_bucket: the name of the S3 bucket where env vars should be + searched. + :param kubernetes_labels: labels to be passed to the kubernetes agent. + :param kubernetes_image: image to use when running through the kubernetes agent. + :param slack_webhooks_file: path to the slack webhooks file within the env + bucket. + """ + self.FLOW_NAME = flow_name + self.FLOW_NAME_UDCS = flow_name.replace("-", "_ ") + self.S3_BUCKET_NAME = ( + DEFAULT_ENV_S3_BUCKET if env_s3_bucket is None else env_s3_bucket + ) + self.SLACK_WEBHOOKS_FILE = ( + DEFAULT_PATH_TO_SLACK_WEBHOOKS_FILE + if slack_webhooks_file is None + else slack_webhooks_file + ) + self.SLACK_WEBHOOKS = None + self.STORAGE = S3( + bucket=self.S3_BUCKET_NAME, + key=DEFAULT_FLOWS_PATH_IN_BUCKET + self.FLOW_NAME + ".py", + stored_as_script=True, + ) + self.KUBERNETES_LABELS = ( + DEFAULT_KUBERNETES_LABELS + if kubernetes_labels is None + else kubernetes_labels + ) + self.KUBERNETES_IMAGE = ( + DEFAULT_KUBERNETES_IMAGE if kubernetes_image is None else kubernetes_image + ) + + self._s3_reader = S3FileReader(s3_client=boto3.client("s3")) + + def fetch_slack_webhooks(self, s3_reader=None) -> None: + """ + Read the slack webhooks file from S3 and store the webhooks in memory. + + :param s3_reader: a client to fetch files from S3. + :return: None + """ + + if s3_reader is None: + s3_reader = self._s3_reader + + self.SLACK_WEBHOOKS = s3_reader.read_json_from_s3_file( + bucket=self.S3_BUCKET_NAME, key=self.SLACK_WEBHOOKS_FILE + ) + + +def build_lolaconfig( + flow_name: str, + env_s3_bucket: str = None, + kubernetes_labels: List = None, + kubernetes_image: str = None, +) -> LolaConfig: + """ + Build a LolaConfig instance from the passed params. + + :param flow_name: the name of the flow. + :param env_s3_bucket: the name of the S3 bucket where env vars should be + searched. + :param kubernetes_labels: labels to be passed to the kubernetes agent. + :param kubernetes_image: image to use when running through the kubernetes agent. + :return: a ready to use LolaConfig instance. + """ + + lolaconfig = LolaConfig( + flow_name=flow_name, + env_s3_bucket=env_s3_bucket, + kubernetes_labels=kubernetes_labels, + kubernetes_image=kubernetes_image, + ) + + lolaconfig.fetch_slack_webhooks() + + return lolaconfig diff --git a/lolafect/utils.py b/lolafect/utils.py new file mode 100644 index 0000000..621e59d --- /dev/null +++ b/lolafect/utils.py @@ -0,0 +1,24 @@ +import json + + +class S3FileReader: + """ + An S3 client along with a few reading utils. + """ + + def __init__(self, s3_client): + self.s3_client = s3_client + + def read_json_from_s3_file(self, bucket: str, key: str) -> dict: + """ + Read a JSON file from an S3 location and return contents as a dict. + + :param bucket: the name of the bucket where the file is stored. + :param key: the path to the file within the bucket. + :return: the file contents. + """ + return json.loads( + self.s3_client.get_object(Bucket=bucket, Key=key)["Body"] + .read() + .decode("utf-8") + ) diff --git a/requirements-dev.txt b/requirements-dev.txt index a4f2f46..86204fb 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,2 +1,4 @@ prefect==1.2.2 -requests==2.28.1 \ No newline at end of file +requests==2.28.1 +boto3==1.26.40 +pytest==7.2.0 \ No newline at end of file diff --git a/setup.py b/setup.py index d1fb455..1b6c89e 100644 --- a/setup.py +++ b/setup.py @@ -23,5 +23,5 @@ setup( package_dir={"lolafect": "lolafect"}, include_package_data=True, python_requires=">=3.7", - install_requires=["prefect==1.2.2", "requests==2.28.1"], + install_requires=["prefect==1.2.2", "requests==2.28.1", "boto3==1.26.40"], ) diff --git a/tests/test_lolaconfig.py b/tests/test_lolaconfig.py new file mode 100644 index 0000000..742de77 --- /dev/null +++ b/tests/test_lolaconfig.py @@ -0,0 +1,35 @@ +from types import SimpleNamespace + +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", + kubernetes_labels=["some_label"], + kubernetes_image="loliloli:latest", + slack_webhooks_file="the_file/is/here.json", + ) + + +def test_lolaconfig_fetches_webhooks_properly(): + + lolaconfig = LolaConfig(flow_name="some-flow") + + fake_s3_reader = SimpleNamespace() + + def mock_read_json_from_s3_file(bucket, key): + return {"a-channel-name": "a-channel-url.com"} + + fake_s3_reader.read_json_from_s3_file = mock_read_json_from_s3_file + + lolaconfig.fetch_slack_webhooks(s3_reader=fake_s3_reader) + + assert type(lolaconfig.SLACK_WEBHOOKS) is dict