corva-worker-python 2.0.0__py3-none-any.whl

worker/partial_rerun_merge/progress.py

@@ -0,0 +1,241 @@
+"""
+This module defines the progress of the partial reruns.
+"""
+
+from dataclasses import asdict, dataclass
+from typing import List, Optional
+
+import requests
+import simplejson as json
+
+from worker.data.api import API
+from worker.data.enums import PartialRerunStatus
+
+
+@dataclass
+class DatasetProgress:
+    """
+    Represents the progress of a dataset.
+
+    Attributes:
+        dataset_id (int): The ID of the dataset.
+        processed_timestamp (int): The timestamp when the dataset was processed.
+        completed (bool): Whether the dataset processing is completed or not.
+        dataset_name (str): The name of the dataset.
+    """
+
+    dataset_id: int
+    processed_timestamp: int
+    completed: bool
+    dataset_name: str
+
+    def is_completed(self) -> bool:
+        """
+        Returns a boolean indicating whether the progress has been completed.
+        """
+        return self.completed
+
+    def is_started(self) -> bool:
+        """
+        Returns a boolean indicating whether the progress has been started.
+        """
+        return self.processed_timestamp > 0
+
+    def mark_completed_at(self, processed_timestamp: int) -> None:
+        """
+        Marks the progress as completed and sets the processed timestamp.
+
+        Args:
+            processed_timestamp (int): The timestamp when the progress was processed.
+        """
+        self.completed = True
+        self.processed_timestamp = processed_timestamp
+
+    @classmethod
+    def set_from_dict(cls, json_object: dict) -> "DatasetProgress":
+        """Set the object from a dict.
+
+        Args:
+            json_object (dict): the input json
+
+        Returns:
+            _type_: DatasetProgress
+        """
+        return cls(**json_object)
+
+    def to_dict(self) -> dict:
+        """
+        Returns a dictionary representation of the Progress object.
+
+        Returns:
+            dict: A dictionary containing the dataset_id, processed_timestamp,
+                    completed, and dataset_name attributes.
+        """
+        return asdict(self)
+
+    def get_trimmed_dataset_name(self) -> str:
+        """
+        Returns the trimmed dataset name by splitting the dataset name
+        with "#" delimiter and returning the last part.
+
+        Example: "corva#circulation.volumetric" -> "circulation.volumetric"
+
+        Returns:
+            str: The trimmed dataset name.
+        """
+        return self.dataset_name.split("#")[-1]
+
+
+class MergingProgress:
+    """
+    Represents the progress of merging a partial rerun for a specific app.
+
+    Attributes:
+        partial_well_rerun_id (int): The ID of the partial rerun.
+        app_id (int): The ID of the app.
+        status (PartialRerunStatus): The status of the partial rerun.
+        dataset_progresses (List[DatasetProgress]): The progress of each dataset
+            in the partial rerun.
+        api: (Worker) The API object used to make requests to the Corva API.
+        is_cache_update_completed (bool): Whether the cache update is completed.
+    """
+
+    def __init__(self, partial_well_rerun_id: int, app_id: int, api: API):
+        """
+        Initializes a new instance of the Progress class.
+
+        Args:
+            partial_rerun_id (int): The ID of the partial rerun.
+            app_id (int): The ID of the application.
+            api: The API object.
+
+        Returns:
+            None
+        """
+        self.partial_well_rerun_id = partial_well_rerun_id
+        self.app_id = app_id
+        self.status = PartialRerunStatus.MERGING
+        self.dataset_progresses: List[DatasetProgress] = []
+
+        self.api = api
+
+        # getting and parsing the progress
+        request = self.api.get(self.get_url_path()).data
+        self.parse_progress(request)
+
+        self.is_cache_update_completed = self.determine_if_cache_update_completed()
+
+        # in case of failure, this will be set to the reason of failure
+        self.fail_reason: Optional[str] = None
+
+    def parse_progress(self, request: dict) -> None:
+        """
+        Parses the progress information retrieved from the API and populates the
+        'dataset_progresses' list.
+
+        Args:
+            request (dict): The progress information retrieved from the API.
+
+        Returns:
+            None
+        """
+        progresses = request.get("included") or []
+
+        for progress in progresses:
+            if progress.get("type") != "partial_well_rerun_dataset_progress":
+                continue
+
+            attributes = progress.get("attributes") or {}
+
+            dataset_progress = DatasetProgress.set_from_dict(attributes)
+            self.dataset_progresses.append(dataset_progress)
+
+    def complete_status(self) -> None:
+        """
+        Marks the partial rerun merge as completed.
+        """
+        self.status = PartialRerunStatus.COMPLETED
+
+    def fail_status(self, reason: str) -> None:
+        """
+        Sets the status of the partial rerun to FAILED and provides a reason for the failure.
+
+        Args:
+            reason (str): The reason for failure.
+        """
+        self.status = PartialRerunStatus.FAILED
+        self.fail_reason = reason
+
+    def get_dataset_progress(self, dataset_name: str) -> Optional[DatasetProgress]:
+        """
+        Gets the progress of the dataset with the given name.
+
+        Args:
+            dataset_name (str): The name of the dataset.
+
+        Returns:
+            DatasetProgress: The progress of the dataset.
+        """
+        return next((dp for dp in self.dataset_progresses if dp.get_trimmed_dataset_name() == dataset_name), None)
+
+    def is_collection_completed(self, collection_name: str) -> bool:
+        """
+        Checks if a collection with the given name has been completed.
+
+        Args:
+            collection_name (str): The name of the collection to check.
+
+        Returns:
+            bool: True if the collection has been completed, False otherwise.
+        """
+        for dataset_progress in self.dataset_progresses:
+            if dataset_progress.get_trimmed_dataset_name() == collection_name:
+                return dataset_progress.completed
+
+        return False
+
+    def determine_if_cache_update_completed(self) -> bool:
+        """
+        Checks if the cache update is completed.
+        Since there is no collection named "cache", it checks if ANY collection is completed,
+        and since cache update needs to be done once, therefore we need to know when the merge
+        will continue in the next iteration, we need to check if the cache update is already completed.
+
+        Returns:
+            bool: True if the cache update has completed for all datasets, False otherwise.
+        """
+        return any(dataset_progress.is_started() for dataset_progress in self.dataset_progresses)
+
+    def get_url_path(self) -> str:
+        """
+        Returns the URL path for the API request.
+
+        Returns:
+            str: URL Path for the API request.
+        """
+        return f"/v2/partial_reruns/{self.partial_well_rerun_id}/app_progress/{self.app_id}"
+
+    def to_dict(self) -> dict:
+        """
+        Converts the merging progress to a dict.
+
+        Returns:
+            dict: progress as a dict.
+        """
+        dataset_progresses = [dataset_progress.to_dict() for dataset_progress in self.dataset_progresses]
+
+        output = dict(app_progress=dict(status=self.status.value, dataset_progresses=dataset_progresses))
+
+        if self.fail_reason:
+            output["app_progress"]["fail_reason"] = self.fail_reason
+
+        return output
+
+    def update_status(self) -> requests.Response:
+        """
+        Sends the updated merging progress to the API.
+        """
+        url_path = self.get_url_path()
+        body = json.dumps(self.to_dict())
+        response = self.api.patch(url_path, data=body)
+        return response.response

worker/state/__init__.py

@@ -0,0 +1,96 @@
+import os
+from typing import List, Union
+
+import redis
+import simplejson as json
+
+from worker.data.json_encoder import JsonEncoder
+from worker.mixins.logging import Logger
+
+
+class RedisHandler:
+    redis_ttl = os.getenv("REDIS_TTL", 90 * 24 * 3600)  # 90 days
+    _redis_connection = None
+
+    @classmethod
+    def _get_redis_connection(cls) -> redis.client.Redis:
+        """
+        Setup redis and get connection if one does not exist already
+        :return:
+        """
+        # Use previous connection if cached
+        if cls._redis_connection:
+            Logger.debug("Using existing connection")
+            return cls._redis_connection
+
+        # If connection does not exist, create a new one
+        cls._redis_connection = cls._new_redis_connection()
+
+        return cls._redis_connection
+
+    @classmethod
+    def _new_redis_connection(cls) -> redis.client.Redis:
+        """
+        Creates a new Redis connection
+        :return:
+        """
+        cache_url = os.getenv("CACHE_URL", None)
+
+        if not cache_url:
+            raise Exception("redis key (CACHE_URL) not found in Environment Variables.")
+
+        _redis_connection = redis.Redis.from_url(cache_url)
+
+        if not _redis_connection:
+            raise Exception(f"Could not connect to Redis with URL: {cache_url}")
+
+        client_id = _redis_connection.client_id()
+        Logger.log(f"Created a new connection with {client_id=}")
+        return _redis_connection
+
+    @classmethod
+    def setup_existing_connection(cls, existing_connection: Union[redis.client.Redis, None]) -> redis.client.Redis:
+        """
+        Set the internal connection variable to existing connection provided in this function
+        The returned output should be saved to a global variable if it is intended to cache the
+        connection across lambda invokes
+        :param existing_connection:
+        :return:
+        """
+        cls._redis_connection = existing_connection
+        return cls._get_redis_connection()
+
+    @classmethod
+    def load_state(cls, state_key: str) -> dict:
+        """
+        Load state from redis
+        :return:
+        """
+        state = cls._get_redis_connection().get(state_key)
+        if state:
+            return json.loads(state)
+
+        return {}
+
+    @classmethod
+    def save_state(cls, state: dict, state_key: str, **kwargs):
+        """
+        Save the state to redis. Uses redis_ttl from kwargs if passed, fallbacks to env var then default 90 days
+        :param state:
+        :param state_key:
+        :return:
+        """
+        # Using redis ttl from kwargs if passed in.
+        redis_ttl = kwargs.pop("redis_ttl", cls.redis_ttl)
+        cls._get_redis_connection().set(
+            state_key, value=json.dumps(state, cls=JsonEncoder, ignore_nan=True), ex=redis_ttl
+        )
+
+    @classmethod
+    def delete_states(cls, state_keys: List[str]):
+        """
+        Delete the states corresponding to given keys
+        :param state_keys: a list of redis keys
+        :return:
+        """
+        cls._get_redis_connection().delete(*state_keys)

worker/state/mixins.py

@@ -0,0 +1,111 @@
+import os
+import sys
+from functools import cached_property
+from typing import List, Optional, Union
+
+import simplejson as json
+
+from worker.data.json_encoder import JsonEncoder
+from worker.mixins.logging import LoggingMixin
+from worker.mixins.rollbar import RollbarMixin
+from worker.state import RedisHandler
+
+
+class RedisMixin(LoggingMixin, RollbarMixin):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.state = {}
+
+    @cached_property
+    def _size_limits(self):
+        state_storage_limit_fatal = os.getenv("STATE_STORAGE_CRITICAL_LIMIT", 1_000.0)  # 1 MB
+        state_storage_limit_warning = os.getenv("STATE_STORAGE_WARNING_LIMIT", 100.0)  # 100 kB
+        size_limits = {"fatal": float(state_storage_limit_fatal), "warning": float(state_storage_limit_warning)}
+        return size_limits
+
+    def load_state(self, state_key: Union[str, None] = None, raise_warnings: bool = False) -> dict:
+        """
+        Load the state from redis
+        :param state_key:
+        :param raise_warnings: Set to True if size limits warnings should be printed
+        :return:
+        """
+        if not state_key:
+            state_key = self.get_formatted_state_key(self.asset_id, self.app_key, self.module_key)
+
+        state = RedisHandler.load_state(state_key)
+        size_object = self.check_state_size(state, state_key, raise_warnings=raise_warnings)
+
+        self.debug(self.asset_id, f"Retrieved state of size {size_object} kb")
+        self.state = state
+        return state
+
+    def save_state(self, state_key: Union[str, None] = None, raise_warnings: bool = True, **kwargs) -> None:
+        """
+        Save the state to redis
+        :param state_key:
+        :param raise_warnings: Set to True if size limits warnings should be printed
+        :return:
+        """
+        if not state_key:
+            state_key = self.get_formatted_state_key(self.asset_id, self.app_key, self.module_key)
+
+        size_object = self.check_state_size(self.state, state_key, raise_warnings=raise_warnings)
+        RedisHandler.save_state(self.state, state_key, **kwargs)
+        self.debug(self.asset_id, f"Saved state of size {size_object} kb")
+
+    def delete_states(self, state_keys: Union[List[str], str]):
+        """
+        Delete state for current module
+        :param state_keys:
+        :return:
+        """
+        if not isinstance(state_keys, list):
+            state_keys = [state_keys]
+
+        RedisHandler.delete_states(state_keys)
+        self.debug(None, "Deleted state from Redis")
+
+    @staticmethod
+    def get_formatted_state_key(asset_id: int, app_key: str, module_key: Optional[str] = None) -> str:
+        """
+        Returns the state key in Corva naming format
+        :param asset_id:
+        :param module_key:
+        :param app_key:
+        :return:
+        """
+        state_key = "corva/{0}.{1}".format(asset_id, app_key)
+
+        if module_key:
+            return "{0}.{1}".format(state_key, module_key)
+
+        return state_key
+
+    def check_state_size(self, state, state_key, raise_warnings=True):
+        """
+        Check the size of the state dictionary and generate warnings if necessary
+        :param state:
+        :param state_key:
+        :param raise_warnings:
+        :return:
+        """
+        size_object = sys.getsizeof(json.dumps(state, cls=JsonEncoder, ignore_nan=True)) / 1024
+
+        if not raise_warnings:
+            return size_object
+
+        size_limit = self._size_limits["fatal"]
+        if size_object > size_limit:
+            message = f"State_key {state_key} is of size {size_object} kb > {size_limit} kb."
+            self.fatal(self.asset_id, message)
+            self.track_message(message, level="critical")
+            return size_object
+
+        size_limit = self._size_limits["warning"]
+        if size_object > size_limit:
+            message = f"State_key {state_key} is of size {size_object} kb > {size_limit} kb."
+            self.warn(self.asset_id, message)
+            return size_object
+
+        return size_object

worker/state/state.py

@@ -0,0 +1,46 @@
+import simplejson as json
+
+from worker.data.json_encoder import JsonEncoder
+
+
+class State(dict):
+    def __init__(self, fields: dict, state: dict = None, *args, **kwargs):
+        """
+        :param fields: is a dictionary that holds the state items. The value of each key is the data type.
+        :param state:
+        :param args:
+        :param kwargs:
+        """
+        self.fields = fields
+
+        super().__init__(*args, **kwargs)
+        if not state:
+            state = {}
+
+        for field, conversion in self.fields.items():
+            self.set_field_if_exists(field, state, conversion)
+
+    def set_field_if_exists(self, field, state, conversion):
+        if not field:
+            return
+
+        value = None
+
+        if field in state and state[field] is not None:
+            value = state[field]
+            if conversion:
+                value = conversion(value)
+
+        self[field] = value
+
+    def to_json(self):
+        output = {field: self.get(field, None) for field in self.fields.keys()}
+        output = json.dumps(output, cls=JsonEncoder, ignore_nan=True)
+
+        return output
+
+    def get(self, key, default=None):
+        value = super().get(key, None)
+        if value is None:
+            value = default
+        return value

worker/test/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa: F401
+
+from worker.test.lambda_function_test_run import AppTestRun

worker/test/lambda_function_test_run.py

@@ -0,0 +1,196 @@
+"""
+This file is used to trigger lambda functions locally and test
+the results of the app on actual data to make sure of the results.
+An example to use this functionality is like the one shows on
+'app_test.py' and then that file can run in command line.
+
+Note: make sure you run your tests on QA environment.
+
+== app_test.py file
+# added the main directory as the first path
+from os.path import dirname
+import sys
+parent = dirname(dirname(__file__))
+sys.path.insert(0, parent)
+
+if __name__ == '__main__':
+    collections = ['collection_to_delete']
+    app = AppTestRun(lambda_function.lambda_handler, collections)
+    app.run()
+
+> python app_test -a 16886 -d True
+"""
+
+import argparse
+from typing import List
+
+from dotenv import load_dotenv
+from tqdm import tqdm
+
+from worker import constants
+from worker.data.operations import (
+    delete_collection_data_of_asset_id,
+    gather_data_for_period,
+    get_one_data_record,
+    point_main_envs,
+)
+from worker.state.mixins import RedisMixin
+from worker.test.utils import create_scheduler_events, get_last_processed_timestamp
+
+load_dotenv(override=True)
+
+
+def strtobool(val):
+    """Convert a string representation of truth to true (1) or false (0).
+
+    True values are 'y', 'yes', 't', 'true', 'on', and '1'.
+    False values are 'n', 'no', 'f', 'false', 'off', and '0'.
+    Raises ValueError if 'val' is anything else.
+    """
+    val = val.lower()
+    if val in ("y", "yes", "t", "true", "on", "1"):
+        return 1
+    elif val in ("n", "no", "f", "false", "off", "0"):
+        return 0
+    else:
+        raise ValueError("invalid truth value %r" % (val,))
+
+
+def generate_runner_parser():
+    """
+    Creating the supporting arguments
+    :return:
+    """
+    parser = argparse.ArgumentParser(description="Run your tests on an asset.")
+    parser.add_argument(
+        "-v",
+        "--environment",
+        "--env",
+        type=str,
+        required=False,
+        help="environment, options: 'local', 'qa', 'staging', 'production'",
+    )
+    parser.add_argument("-a", "--asset_id", "--id", type=int, required=True, help="set asset_id")
+    parser.add_argument(
+        "-s", "--start_timestamp", "--start", type=int, required=False, default=None, help="start timestamp"
+    )
+    parser.add_argument("-e", "--end_timestamp", "--end", type=int, required=False, default=None, help="end timestamp")
+    parser.add_argument(
+        "-i",
+        "--timestep",
+        "--step",
+        type=int,
+        required=False,
+        default=60,
+        help="trigger the lambda function once every step",
+    )
+    parser.add_argument(
+        "-d",
+        "--to_delete",
+        "--delete",
+        type=strtobool,
+        required=False,
+        default=False,
+        help="to delete the state and data",
+    )
+    return parser
+
+
+class AppTestRun:
+    def __init__(self, lambda_handler, collections: List[str], args: List[str] = None):
+        """
+        :param lambda_handler: lambda handler function to run
+        :param collections: collections to erase in case to_delete is on
+        :param constants: constants module of the whole app, if there is only one app in the whole repository
+        pass nothing and the code takes care of that
+        :param args: arguments for the run
+        """
+        self.lambda_handler = lambda_handler
+        self.collections = collections
+
+        self.event_type = None
+        self.progress = None
+
+        if args is None:
+            parser = generate_runner_parser()
+            args = parser.parse_args()
+
+        self.initialize(args)
+
+    def initialize(self, args):
+        environment = args.environment
+
+        # directing main environment variables to the provided env: API_KEY, API_ROOT_URL, CACHE_URL
+        point_main_envs(environment)
+
+        asset_id = args.asset_id
+
+        start_timestamp = args.start_timestamp
+        end_timestamp = args.end_timestamp
+        step = args.timestep
+        self.event_type = constants.get("global.event-type")
+        to_delete = args.to_delete
+
+        state_keys = construct_state_keys(asset_id)
+
+        if not start_timestamp:
+            start_timestamp = get_one_data_record(asset_id, timestamp_sort=+1).get("timestamp")
+
+        if not end_timestamp:
+            end_timestamp = get_one_data_record(asset_id, timestamp_sort=-1).get("timestamp")
+
+        if to_delete:
+            delete_state_data_of_asset_id(state_keys)
+            delete_collection_data_of_asset_id(asset_id, self.collections)
+            print("Deleted relevant Redis states and collections for this asset!")
+
+        start_timestamp = get_last_processed_timestamp(asset_id, state_keys[0]) or start_timestamp
+        print(f"asset_id: {asset_id}, timestamp interval: [{start_timestamp}, {end_timestamp}]")
+
+        events = create_scheduler_events(asset_id, start_timestamp, end_timestamp, step)
+
+        self.progress = tqdm(events, ncols=150)
+
+    def run(self):
+        print("\nRunning the main module started ...")
+        for event in self.progress:
+            schedule_time = str(int(event[0][0]["schedule_start"] / 1000))
+            if self.event_type == "wits_stream":
+                wits = gather_data_for_period(
+                    int(event[0][0]["asset_id"]),
+                    int(event[0][0]["schedule_start"] / 1000),
+                    int(event[0][0]["schedule_end"] / 1000),
+                )
+                if not wits:
+                    continue
+
+                event = [{"metadata": {}, "records": wits}]
+
+            self.lambda_handler(event, None)
+            self.progress.set_description(schedule_time)
+
+
+def construct_state_keys(asset_id: int) -> List[str]:
+    """
+    Get the constants dict of an app and construct the storage keys.
+    Note: the first key is the global app key.
+    :param asset_id: well asset id
+    :return: a list of storage keys
+    """
+    app_key = constants.get("global.app-key")
+    module_keys = list(constants.get(app_key).keys())
+    state_app = RedisMixin()
+    global_state_key = state_app.get_formatted_state_key(asset_id, app_key)
+
+    module_state_keys = [state_app.get_formatted_state_key(asset_id, app_key, module) for module in module_keys]
+
+    return [global_state_key, *module_state_keys]
+
+
+def delete_state_data_of_asset_id(state_keys: List[str]):
+    """
+    Delete the state from given storage type
+    :param state_keys:
+    :return:
+    """
+    RedisMixin().delete_states(state_keys)