Merge pull request #1 from lolamarket/feature/lolaconfig

Feature/lolaconfig

README.md

@@ -1,3 +1,13 @@
 # Lolafect
 
 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.

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/"

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

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")
+        )

requirements-dev.txt

@@ -1,2 +1,4 @@
 prefect==1.2.2
 requests==2.28.1
+boto3==1.26.40
+pytest==7.2.0

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"],
 )

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