115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
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,
|
|
DEFAULT_ENV_FILE_PATH
|
|
)
|
|
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,
|
|
env_file_path: 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 env_file_path: the path to the environment file.
|
|
: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.ENV_FILE_PATH = (
|
|
DEFAULT_ENV_FILE_PATH if env_file_path is None else env_file_path
|
|
)
|
|
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
|