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

worker/data/task_handler.py

@@ -0,0 +1,82 @@
+import os
+
+import simplejson as json
+
+from worker import API
+from worker.data.enums import LambdaStates
+from worker.data.json_encoder import JsonEncoder
+from worker.mixins.logging import LoggingMixin
+from worker.mixins.rollbar import RollbarMixin
+
+
+class TaskHandler(LoggingMixin, RollbarMixin):
+    def __init__(self, *args, **kwargs) -> None:
+        self._api = kwargs.pop("api", API())
+        super().__init__(*args, **kwargs)
+        self.event = {}
+        self.task_id = ""
+        self.task = {}
+
+    def process(self, *args, **kwargs):
+        """
+        Provides a placeholder for actual logic code for each app.
+        e.g. Instantiate the app and run it
+        This method must be overriden by each app
+        :return:
+        """
+        raise NotImplementedError("process method must be defined before instantiating this class")
+
+    def run_process(self, event: dict, *args, **kwargs):
+        """
+        Runs the task handler workflow.
+        This method would be called after instantiation of the TaskHandler class
+        :param event:
+        :return:
+        """
+        self.event = event
+        self.task_id = self.event.get("task_id")
+        self.get_task(self.task_id)
+
+        # AWS_EXECUTION_ENV is a reserved environment variable it will hold the
+        # runtime environment eg python3.6, java11 etc.
+        # We do not have to set it on our lambda functions, it will be set by AWS
+        # We can use this to check if we are running on hosted or local
+        is_hosted = os.environ.get("AWS_EXECUTION_ENV", False)
+
+        test_mode = self.task.get("properties", {}).get("test_mode", False)
+
+        if test_mode and is_hosted:
+            asset_id = self.task.get("asset_id", 0)
+            self.warn(asset_id, f"Not processing on AWS due to test mode. task_id: {self.task_id} maybe running.")
+            return
+
+        try:
+            data = self.process(self.event, *args, **kwargs)
+            request_body = {"task": {"payload": {"data": data, "message": LambdaStates.SUCCEEDED.value}}}
+            self.update_task(self.task_id, "success", request_body)
+            return data
+
+        except Exception as exception:
+            request_body = {"task": {"fail_reason": f"Message: {LambdaStates.FAILED.value}, Exception: {exception}"}}
+            self.update_task(self.task_id, "fail", request_body)
+            raise exception
+
+    def get_task(self, task_id: str) -> None:
+        """
+        Gets the task from API endpoint for the given task_id
+        :param task_id:
+        :return:
+        """
+        path = "/v2/tasks/{0}".format(task_id)
+        self.task = self._api.get(path).data
+
+    def update_task(self, task_id: str, status: str, request_body: dict) -> None:
+        """
+        Updates the task status on the API endpoint in case of success or failure
+        :param task_id:
+        :param status: [success, fail]
+        :param request_body:
+        :return:
+        """
+        path = "/v2/tasks/{0}/{1}".format(task_id, status)
+        self.task = self._api.put(path, data=json.dumps(request_body, cls=JsonEncoder, ignore_nan=True)).data

worker/data/two_way_dict.py

@@ -0,0 +1,17 @@
+class TwoWayDict(dict):
+    def __setitem__(self, key, value):
+        # Remove any previous connections with these values
+        if key in self:
+            del self[key]
+        if value in self:
+            del self[value]
+        dict.__setitem__(self, key, value)
+        dict.__setitem__(self, value, key)
+
+    def __delitem__(self, key):
+        dict.__delitem__(self, self[key])
+        dict.__delitem__(self, key)
+
+    def __len__(self):
+        """Returns the number of connections"""
+        return dict.__len__(self) // 2

worker/data/unit_conversions.py

@@ -0,0 +1,5 @@
+# unit conversions
+KG_M3_to_PPG = 0.00834540445319611
+GPM_SEC_to_FT3 = (1 / 60) * 0.133681
+US_LIQUID_GAL_to_FOOT3 = 0.133681
+US_LIQUID_GAL_to_INCH3 = 231

worker/data/wits.py

@@ -0,0 +1,323 @@
+from typing import Dict, Union
+
+import numpy as np
+
+from worker.data import operations
+from worker.data.activity import Activity
+from worker.data.enums import ChannelStatus, DataStatus
+
+
+class Channel(float):
+    def __new__(cls, value=None, status=None):
+        # If only a valid value is passed, data status is automatically set to VALID
+        # If None is passed as value, then data status is automatically set to MISSING
+        # If both value and data status are passed, then it uses the passed data status
+        if value is None:
+            return None
+
+        if not isinstance(value, float):
+            try:
+                value = float(value)
+            except ValueError or TypeError:
+                return None
+
+        if value is not None and not status:
+            status = DataStatus.VALID
+
+        ch = float.__new__(cls, value)
+        ch._status = status
+        return ch
+
+    @property
+    def status(self) -> DataStatus:
+        """
+        - Status can be valid of missing
+        - Data status is only available for non None wits attributes
+        - Missing represents a value that has been forward filled or interpolated
+
+        :return: DataStatus
+        """
+        return self._status
+
+
+class WITS:
+    """
+    - Create a wits object by WITS.set_wits(dataset: Union[list, dict])
+    - Access any wits data attribute by wits.hole_depth
+    - Access any data status of a wits data attribute by wits.hole_depth.status
+    - Data status is only available for non None wits attributes
+    """
+
+    data = {
+        "entry_at": int,
+        "hole_depth": Channel,
+        "bit_depth": Channel,
+        "block_height": Channel,
+        "hook_load": Channel,
+        "weight_on_bit": Channel,
+        "rop": Channel,
+        "rotary_rpm": Channel,
+        "rotary_torque": Channel,
+        "mud_flow_in": Channel,
+        "mud_flow_out_percent": Channel,
+        "standpipe_pressure": Channel,
+        "diff_press": Channel,
+        "pump_spm_1": Channel,
+        "pump_spm_2": Channel,
+        "pump_spm_3": Channel,
+        "pump_spm_4": Channel,
+        "pump_spm_total": Channel,
+        "trip_tank_volume_1": Channel,
+        "trip_tank_volume_2": Channel,
+        "active_pit_volume": Channel,
+        "lag_depth": Channel,
+        "state": Activity,
+        "boost_pump_flow_in": Channel,
+        "annular_back_pressure": Channel,
+    }
+    metadata = {"drillstring": str}
+    spm_channels = ["pump_spm_1", "pump_spm_2", "pump_spm_3", "pump_spm_4"]
+    status = {}
+
+    def __init__(self, wits_json: dict = None, **kwargs):
+        if not wits_json:
+            wits_json = {}
+
+        # Serialize functionality in Worker unpacks all the attributes. If that's how the class is called, repack.
+        if kwargs:
+            wits_json.update(dict(kwargs))
+            wits_json.pop("_wits_", None)
+
+        self.wits_json = wits_json
+
+        data = wits_json.get("data", {})
+        data_raw = wits_json.get("data_raw", {})
+
+        metadata = wits_json.get("metadata", {})
+        status = wits_json.get("status", {})
+
+        for key in self.data.keys():
+            # If type is not Channel or if the value is None, just setattr
+            if self.data.get(key) != Channel or data.get(key, None) is None:
+                setattr(self, key, data.get(key, None))
+                continue
+
+            key_data_status = DataStatus.VALID.value
+            # If key is in data_raw, there are 2 possibilities - Missing (Value is None or -999.25) or Overridden
+            if key in data_raw:
+                key_data_status = DataStatus.OVERRIDDEN.value
+                # If key in data_raw and value is null (None or -999.25), then we just setattr None
+                if operations.is_null(data_raw.get(key, None)):
+                    setattr(self, key, data.get(key, None))
+                    continue
+
+            # If type is Channel and (key not in data_raw or (key in data_raw and value is not null))
+            setattr(self, key, Channel(data.get(key), DataStatus(data.get(f"{key}_data_status", key_data_status))))
+
+        for key in self.metadata.keys():
+            setattr(self, key, metadata.get(key, None))
+        for key in self.status.keys():
+            setattr(self, key, status.get(key, None))
+
+        setattr(self, "timestamp", wits_json.get("data").get("entry_at"))
+        setattr(self, "timestamp_read", wits_json.get("timestamp_read"))
+
+        if getattr(self, "pump_spm_total") is None:
+            self.set_pump_spm_total()
+
+        setattr(self, "total_trip_tank_volume", self.get_total_trip_tank_volume())
+
+    def __getattr__(self, item):
+        return None
+
+    def __setattr__(self, key, value):
+        """
+        When trying to set an attribute it automatically tries to typecast the value.
+        :param key: key
+        :param value: value
+        :return:
+        """
+        if key not in {**self.data, **self.status}.keys():
+            self.__dict__[key] = value
+            return
+
+        data_type = {**self.data, **self.status}.get(key)
+        try:
+            if not isinstance(value, Channel):
+                value = data_type(value)
+            self.__dict__[key] = value
+        except TypeError:
+            self.__dict__[key] = None
+        except ValueError:
+            self.__dict__[key] = None
+
+    def __iter__(self):
+        yield self
+
+    @classmethod
+    def set_wits(cls, dataset: Union[list, dict]):
+        """
+        Class method to create WITS objects using a dict or a list of dicts
+
+        :param dataset: wits record or a list of wits records
+        :return: WITS object or a list of WITS objects
+        """
+        if isinstance(dataset, dict):
+            return cls(dataset)
+
+        if isinstance(dataset, list):
+            return [cls(record) for record in dataset]
+
+    def get_as_dict(self) -> [None, dict]:
+        """
+        Get the wits record as a dictionary
+
+        :return: wits record as a dict
+        """
+        if not self.wits_json:
+            return None
+
+        self.wits_json.pop("data_raw", None)
+        self.wits_json["data"] = {key: getattr(self, key) for key in self.data}
+        self.wits_json["data"].update(
+            {
+                f"{key}_data_status": getattr(self, key).status.value
+                for key in self.data
+                if self.data.get(key) == Channel and getattr(self, key) is not None
+            }
+        )
+        # convert state from Activity type to str
+        self.wits_json["data"]["state"] = self.wits_json["data"]["state"].value
+        self.wits_json["status"] = {key: getattr(self, key) for key in self.status}
+
+        return self.wits_json
+
+    def get_wits_data(self, columns=None) -> dict:
+        """
+        Get the data part of the wits record as a dictionary
+
+        :param columns: Specify if you only want a select list of columns
+        :return: data part of wits record as a dict
+        """
+        if not columns:
+            data = {key: getattr(self, key) for key in self.data}
+            return data
+
+        data = {key: getattr(self, key) for key in columns}
+        return data
+
+    def spm_status(self, low_spm_threshold: float) -> [ChannelStatus, None]:
+        """
+        Get spm status. Compares total of all SPM channels against SPM threshold.
+
+        :param low_spm_threshold: threshold
+        :return: None if channels are unavailable, OFF if less than threshold and ON if greater
+        """
+
+        # If all SPM channels are empty, return None
+        if not any(operations.is_number(getattr(self, spm_channel)) for spm_channel in self.spm_channels):
+            return ChannelStatus.MISSING
+
+        if getattr(self, "pump_spm_total") > low_spm_threshold:
+            return ChannelStatus.ON
+
+        return ChannelStatus.OFF
+
+    def mud_flow_in_status(self, low_mud_flow_in_threshold: float) -> [ChannelStatus, None]:
+        """
+        Get flow rate status. Compares mud flow in channel against flow rate threshold.
+
+        :param low_mud_flow_in_threshold: threshold
+        :return: None if channel is unavailable, OFF if less than threshold and ON if greater
+        """
+        if not operations.is_number(self.mud_flow_in):
+            return ChannelStatus.MISSING
+
+        if self.mud_flow_in > low_mud_flow_in_threshold:
+            return ChannelStatus.ON
+
+        return ChannelStatus.OFF
+
+    def check_channel_availability(self, required_wits_channels: list) -> bool:
+        """
+        This function checks if all the required channels are available in the wits record
+
+        :param required_wits_channels: List of required channels
+        :return: Boolean value. True if all the required channels are available, else False
+        """
+        return all(
+            operations.is_in_and_not_none(self.get_wits_data(required_wits_channels), channel)
+            for channel in required_wits_channels
+        )
+
+    def set_pump_spm_total(self):
+        """This function sets the pump_spm_total attribute value based on individual spm values"""
+
+        pump_spm_total = sum(
+            getattr(self, spm_channel)
+            for spm_channel in self.spm_channels
+            if (getattr(self, spm_channel) and getattr(self, spm_channel) >= 0)
+        )
+
+        setattr(self, "pump_spm_total", pump_spm_total)
+
+    def get_total_trip_tank_volume(self) -> float:
+        """
+        This function returns the total trip tank volume if more than one trip tanks are available
+
+        :return: Total trip tank volume
+        """
+        ttk1 = self.trip_tank_volume_1 or 0
+        ttk2 = self.trip_tank_volume_2 or 0
+
+        return ttk1 + ttk2
+
+    def validate_range(self, channel_range_map: Dict[str, Dict]):
+        """
+        This method can be used to validate the range of the given channels and apply the proper operation on the data
+        :param channel_range_map: a dictionary of channel to operation and range ([min, max]) mapping
+        example:
+            {
+                "mud_flow_out_percent": {"range": [0, 100], "operation": "clip"}
+            }
+            options for 'operation': ['clip', None]
+        :return: None,
+        """
+        for channel, properties in channel_range_map.items():
+            orig_value = getattr(self, channel, None)
+            if orig_value is None:
+                continue
+
+            valid_range = properties.get("range")
+            operation = properties.get("operation", None)
+
+            new_value = None
+
+            if operation == "clip":
+                new_value = np.clip(orig_value, *valid_range)
+            elif operation is None:
+                if valid_range[0] <= orig_value <= valid_range[-1]:
+                    new_value = orig_value
+            else:
+                raise ValueError("Invalid operation provided!")
+
+            setattr(self, channel, new_value)
+
+
+def serialize(obj: Union[list, WITS]) -> Union[list, dict, None]:
+    """
+    This function is used to serialize a WITS object or a list of WITS objects.
+    Use WITS.set_wits(dataset: Union[list, dict]) to deserialize.
+
+    :param obj: WITS object(s) to serialize.
+    :return: return a serialized object
+    """
+    if isinstance(obj, list):
+        s_obj = [each.get_as_dict() for each in obj]
+        return s_obj
+
+    if isinstance(obj, WITS):
+        s_obj = obj.get_as_dict()
+        return s_obj
+
+    return None

worker/event/__init__.py

@@ -0,0 +1,53 @@
+import abc
+from functools import cached_property
+
+import simplejson as json
+
+from worker import constants
+from worker.data.enums import EventType
+from worker.data.json_encoder import JsonEncoder
+
+
+class Event(abc.ABC):
+    """An event class that holds the events of a single asset_id."""
+
+    def __init__(self, event_type: EventType):
+        self.event_type: EventType = event_type
+
+        self.asset_id: int = None
+        self.records = []
+        self.app_connection_id = None
+
+    def add_records(self, new_records: list) -> None:
+        self.records.extend(new_records)
+
+    @abc.abstractmethod
+    def complete_event(self, api) -> None:
+        raise NotImplementedError
+
+    def __len__(self):
+        return len(self.records)
+
+    def __getitem__(self, index: int):
+        return self.records[index]
+
+    @cached_property
+    def is_posting_to_message_producer(self) -> bool:
+        """
+        Whether to post to message producer for the wits stream and scheduler events or not.
+        If there is no lambda app following your app then this should be false,
+        because this process is intracting with the API which slows the process down.
+        """
+        is_posting: bool = constants.get("global.post-to-message-producer", False)
+        return is_posting
+
+    @abc.abstractmethod
+    def build_message_producer_payload(self) -> dict:
+        raise NotImplementedError
+
+    def post_to_message_producer(self, api):
+        if not self.is_posting_to_message_producer or not self.app_connection_id:
+            return
+
+        payload = self.build_message_producer_payload()
+        api.post(path="/v1/message_producer", data=json.dumps(payload, cls=JsonEncoder, ignore_nan=True))

worker/event/event_handler.py

@@ -0,0 +1,90 @@
+import itertools
+from typing import Dict, Union
+
+from worker.app import App
+from worker.data.api import API
+from worker.data.enums import EventType
+from worker.data.operations import get_cleaned_event_and_type
+from worker.event import Event
+from worker.event.scheduled import ScheduledEvent, SingleScheduledEvent
+from worker.event.stream import StreamEvent
+from worker.exceptions import EventFormatError
+from worker.partial_rerun_merge.merge import PartialRerunMerge
+
+api = API()
+
+
+class EventHandler:
+    def __init__(self, app: App, merger: PartialRerunMerge):
+        self.app: App = app
+        self.event_type: EventType = None
+        self.event_by_asset_id: Dict[int, Event] = None
+        self.merger: PartialRerunMerge = merger
+
+    def process(self, event: Union[str, dict]):
+        """
+        The whole process that is performed on an event including the
+        loading, handling state, running, and completing the event.
+
+        Args:
+            event (Union[str, dict]): lambda handler event
+        """
+        self._load(event)
+        self._run()
+
+    def _load(self, event: Union[str, dict]):
+        """
+        Cleaning the event and group events based on their asset ids
+
+        Args:
+            event (Union[str, dict]): lambda handler event
+        """
+        event, event_type = get_cleaned_event_and_type(event)
+        self.event_type = event_type
+        self.event_by_asset_id = self.format_event(self.event_type, event)
+
+    def _run(self):
+        """Full run of the events based on their asset id"""
+        for _asset_id, event in self.event_by_asset_id.items():
+            if self.event_type == EventType.PARTIAL_RERUN:
+                self.merger.perform_merge(event.get("data", {}))
+                continue
+
+            self.app.load(self.event_type, event)
+            self.app.run_modules()
+            self.app.save_state()
+            event.complete_event(api)
+
+    @staticmethod
+    def format_event(event_type: EventType, event: Union[list, dict]) -> dict:
+        """
+        validate the wits_stream event, flatten and organize the data into a desired format
+        :param event_type: type of event
+        :param event: the wits or scheduler json event
+        :return: a dict of records that are grouped by the asset_ids
+        """
+        if event_type == EventType.PARTIAL_RERUN:
+            if not isinstance(event, dict):
+                raise EventFormatError("Invalid event!")
+
+            return {event["data"]["asset_id"]: event.copy()}
+
+        elif event_type == EventType.SCHEDULER:
+            if not isinstance(event[0], list):
+                raise EventFormatError("Invalid event!")
+
+            # Scheduler events type is 'list of lists'; flattening into a single list
+            events = [SingleScheduledEvent(item) for sublist in event for item in sublist]
+            merging_function = ScheduledEvent
+
+        else:  # 'wits_stream'
+            events = [StreamEvent(each) for each in event]
+            merging_function = StreamEvent.merge
+
+        # sorting is required otherwise we only capture the last group of each asset_id
+        events.sort(key=lambda single_event: single_event.asset_id)
+        groups = itertools.groupby(events, key=lambda single_event: single_event.asset_id)
+
+        events_by_asset_id = {group: merging_function(list(dataset)) for group, dataset in groups}
+
+        return events_by_asset_id

worker/event/scheduled.py

@@ -0,0 +1,64 @@
+from typing import List, Union
+
+from worker.event import Event, EventType
+
+
+class SingleScheduledEvent:
+    def __init__(self, single_stream_event: dict):
+        self.asset_id = single_stream_event.get("asset_id") or 0
+
+        self.start_time = single_stream_event.get("schedule_start", 0) / 1000
+        self.end_time = single_stream_event.get("schedule_end", 0) / 1000
+        self.schedule_id = single_stream_event.get("schedule")
+        self.app_connection_id = single_stream_event.get("app_connection")
+
+    def complete_event(self, api) -> None:
+        """Sets schedule as completed."""
+        if not self.schedule_id:
+            return
+
+        api.post(path=f"/scheduler/{self.schedule_id}/completed")
+
+
+class ScheduledEvent(Event):
+    """A scheduled event class that holds the events of a single asset_id."""
+
+    def __init__(self, new_records: Union[SingleScheduledEvent, List[SingleScheduledEvent]]):
+        super().__init__(EventType.SCHEDULER)
+
+        if new_records:
+            self.add_records(new_records)
+
+        if self.records:
+            self.app_connection_id = self.records[-1].app_connection_id
+
+    def add_records(self, new_records: Union[SingleScheduledEvent, List[SingleScheduledEvent]]):
+        if isinstance(new_records, SingleScheduledEvent):
+            new_records = [new_records]
+
+        if not self.records:
+            self.asset_id = new_records[0].asset_id
+
+        super().add_records(new_records)
+
+    def complete_event(self, api) -> None:
+        """
+        Two steps can happen for scheduler events.
+        1. Sets schedule as completed; this happens for all the events.
+        2. If another lambda function is following your lambda then it should
+           send a message to message producer as well.
+
+        Args:
+            api (API):
+        """
+        # Step 1: mark all the events as completed.
+        for each in self.records:
+            each.complete_event(api)
+
+        # Step 2: in case another lambda function following your lambda
+        self.post_to_message_producer(api)
+
+    def build_message_producer_payload(self) -> dict:
+        data = [{"timestamp": self.records[-1].start_time}]
+
+        return {"app_connection_id": self.app_connection_id, "asset_id": self.asset_id, "data": data}

worker/event/stream.py

@@ -0,0 +1,48 @@
+from typing import List, Union
+
+from worker import constants
+from worker.data.operations import get_data_by_path
+from worker.event import Event, EventType
+
+
+class StreamEvent(Event):
+    """A stream event class that holds the events of a single asset_id."""
+
+    def __init__(self, event: dict, is_posting_to_message_producer: bool = False):
+        super().__init__(EventType.STREAM)
+
+        self.metadata = event.get("metadata") or {}
+        self.records = event.get("records") or []
+
+        self.asset_id = 0
+        if self.records:
+            self.asset_id = self.records[0].get("asset_id")
+
+        app_key = constants.get("global.stream_app_key")
+        self.app_connection_id = get_data_by_path(
+            self.metadata, f"apps.{app_key}.app_connection_id", func=int, default=None
+        )
+
+    @classmethod
+    def merge(cls, events: List["StreamEvent"]) -> Union["StreamEvent", None]:
+        """Merge stream events of the same asset id into one"""
+        if not events:
+            return None
+
+        merged_event = events[0]
+        for each in events[1:]:
+            merged_event.add(each)
+
+        return merged_event
+
+    def add(self, other: "StreamEvent"):
+        if self.asset_id != other.asset_id:
+            raise Exception(f"Events of different assets can not be merged; {self.asset_id} and {other.asset_id}!")
+
+        self.add_records(other.records)
+
+    def build_message_producer_payload(self) -> dict:
+        return {"app_connection_id": self.app_connection_id, "asset_id": self.asset_id, "data": self.records}
+
+    def complete_event(self, api) -> None:
+        self.post_to_message_producer(api)