corva-worker-python 1.1.1__tar.gz → 1.2.1__tar.gz

{corva-worker-python-1.1.1 → corva-worker-python-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.2
 Name: corva-worker-python
-Version: 1.1.1
+Version: 1.2.1
 Summary: SDK for interacting with Corva
 Home-page: https://github.com/corva-ai/corva-worker-python
 Author: Jordan Ambra <jordan.ambra@corva.ai>, Mohammadreza Kamyab <m.kamyab@corva.ai>

{corva-worker-python-1.1.1 → corva-worker-python-1.2.1}/corva_worker_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.2
 Name: corva-worker-python
-Version: 1.1.1
+Version: 1.2.1
 Summary: SDK for interacting with Corva
 Home-page: https://github.com/corva-ai/corva-worker-python
 Author: Jordan Ambra <jordan.ambra@corva.ai>, Mohammadreza Kamyab <m.kamyab@corva.ai>

{corva-worker-python-1.1.1 → corva-worker-python-1.2.1}/setup.py

@@ -18,7 +18,7 @@ setup(
     name='corva-worker-python',
     author='Jordan Ambra <jordan.ambra@corva.ai>, Mohammadreza Kamyab <m.kamyab@corva.ai>',
     url='https://github.com/corva-ai/corva-worker-python',
-    version='1.1.1',
+    version='1.2.1',
     classifiers=classifiers,
     python_requires='>=3.8',
     description='SDK for interacting with Corva',

{corva-worker-python-1.1.1 → corva-worker-python-1.2.1}/testcase/app/app_lambda.py

@@ -20,7 +20,7 @@ def lambda_handler(event, context):
 
     app = DrillingEfficiency()
 
-    merger = AppPartialRerunMerger(app, API(), context, Logger.get_logger())
+    merger = AppPartialRerunMerger(app, API(), Logger.get_logger())
 
     event_handler = EventHandler(app, merger)
     event_handler.process(event)

{corva-worker-python-1.1.1 → corva-worker-python-1.2.1}/testcase/app/app_partial_rerun_merger.py

@@ -10,10 +10,10 @@ SCHEMA = MergingSchemaModel(collections=collections, modules=None)
 
 
 class AppPartialRerunMerger(PartialRerunMerge):
-    def __init__(self, app, api, context, logger):
+    def __init__(self, app, api, logger):
         # adding the required modules to the schema
         schema = SCHEMA
         # app itself has a cache so it needs to be added to the schema as well
         schema.modules = [app] + app.get_active_modules()
 
-        super().__init__(schema, api, context, logger)
+        super().__init__(schema, api, logger)

{corva-worker-python-1.1.1 → corva-worker-python-1.2.1}/worker/event/event_handler.py

@@ -47,7 +47,7 @@ class EventHandler:
         """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)
+                self.merger.perform_merge(event.get("data", {}))
                 continue
 
             self.app.load(self.event_type, event)

{corva-worker-python-1.1.1 → corva-worker-python-1.2.1}/worker/partial_rerun_merge/merge.py

@@ -37,7 +37,6 @@ class PartialRerunMerge:
         self,
         schema: MergingSchemaModel,
         api: API,
-        context,
         logger,
     ):
         """
@@ -46,7 +45,6 @@ class PartialRerunMerge:
         self.schema = schema
 
         self.api = api
-        self.lambda_context = context
         self.logger = logger
 
         remaining_seconds_threshold = os.getenv("PARTIAL_RERUN_REMAINING_SECONDS_THRESHOLD")
@@ -61,6 +59,7 @@ class PartialRerunMerge:
         self.start_timestamp: Optional[int] = None
         self.end_timestamp: Optional[int] = None
         self.rerun_mode: Optional[RerunMode] = None
+        self.run_until: Optional[int] = None
 
         self.start_hole_depth: Optional[float] = None
         self.end_hole_depth: Optional[float] = None
@@ -72,6 +71,10 @@ class PartialRerunMerge:
         Performs a merge operation by updating the cache state,
         merging collections, and updating the status.
 
+        :param event: the event which should be the data part of
+            a an event dict, or a python sdk partial rerun event object
+            that is converted to a dict
+
         :raises Exception: if an error occurs during the merge operation
         """
         self.preprocess(event)
@@ -88,20 +91,23 @@ class PartialRerunMerge:
     def preprocess(self, event: dict) -> None:
         """
         Performs any necessary preprocessing before the merge operation.
-        """
-        # Add your preprocessing logic here
-        event = event.get("data") or {}
 
-        self.partial_well_rerun_id = event.get("partial_well_rerun_id")
-        self.app_id = event.get("app_id")
-        self.original_asset_id = event.get("asset_id")
-        self.rerun_asset_id = event.get("rerun_asset_id")
+        :param event: the event which should be the data part of
+            a an event dict, or a python sdk partial rerun event object
+            that is converted to a dict
+        """
+        self.partial_well_rerun_id = event["partial_well_rerun_id"]
+        self.app_id = event["app_id"]
+        self.original_asset_id = event["asset_id"]
+        self.rerun_asset_id = event["rerun_asset_id"]
 
         segment = event.get("source_type")
         if segment != "drilling":
             raise ValueError(f"Invalid source type: {segment}")
 
-        self.start_timestamp = event.get("start")
+        self.run_until = event["run_until"]
+
+        self.start_timestamp = event["start"]
         start_wits = self.get_wits_at_or(self.rerun_asset_id, self.start_timestamp, "after")
 
         if end := event.get("end"):
@@ -111,10 +117,13 @@ class PartialRerunMerge:
             end_wits = self.get_wits_at_or(self.rerun_asset_id, self.MAX_TIMESTAMP, "before")
             self.end_timestamp = end_wits.get("timestamp")
 
-        self.start_hole_depth = start_wits.get("data").get("hole_depth")
-        self.end_hole_depth = end_wits.get("data").get("hole_depth")
+        if not self.end_timestamp:
+            self.end_timestamp = self.MAX_TIMESTAMP
+
+        self.start_hole_depth = start_wits.get("data", {}).get("hole_depth")
+        self.end_hole_depth = end_wits.get("data", {}).get("hole_depth")
 
-        self.rerun_mode = RerunMode(event.get("rerun_mode"))
+        self.rerun_mode = RerunMode(event["rerun_mode"])
 
         self.merging_progress = MergingProgress(self.partial_well_rerun_id, self.app_id, self.api)
 
@@ -267,19 +276,19 @@ class PartialRerunMerge:
         :param apply_half: whether to apply half of the remaining seconds threshold
         :return: True if there is enough time, False otherwise
         """
-        if not self.lambda_context:
+        if not self.run_until:
             return True
 
-        remaining_seconds = self.get_remaining_seconds()
+        remaining_seconds = self._get_remaining_seconds()
         threshold = self.REMAINING_SECONDS_THRESHOLD_HALF if apply_half else self.REMAINING_SECONDS_THRESHOLD
         return remaining_seconds > threshold
 
-    def get_remaining_seconds(self) -> int:
+    def _get_remaining_seconds(self) -> float:
         """
         Gets the remaining seconds before the Lambda function times out.
         :return: the remaining seconds
         """
-        return self.lambda_context.get_remaining_time_in_millis() // 1000
+        return self.run_until - time.time()
 
     def get_wits_at_or(self, asset_id: int, timestamp: Optional[int], direction: Literal["before", "after"]) -> dict:
         """

{corva-worker-python-1.1.1 → corva-worker-python-1.2.1}/worker/partial_rerun_merge/progress.py

@@ -130,7 +130,7 @@ class MergingProgress:
     def parse_progress(self, request: dict) -> None:
         """
         Parses the progress information retrieved from the API and populates the
-        'datasetProgresses' list.
+        'dataset_progresses' list.
 
         Args:
             request (dict): The progress information retrieved from the API.
@@ -138,10 +138,13 @@ class MergingProgress:
         Returns:
             None
         """
-        collections = request["included"]
+        collections = request.get("included") or []
 
         for collection in collections:
-            if not (attributes := collection.get("attributes")):
+            attributes = collection.get("attributes") or {}
+            dataset_name = attributes.get("dataset_name")
+
+            if not dataset_name:
                 continue
 
             dataset_progress = DatasetProgress.set_from_dict(attributes)