omnata-plugin-runtime 0.11.7a324__tar.gz → 0.11.7a325__tar.gz

{omnata_plugin_runtime-0.11.7a324 → omnata_plugin_runtime-0.11.7a325}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: omnata-plugin-runtime
-Version: 0.11.7a324
+Version: 0.11.7a325
 Summary: Classes and common runtime components for building and running Omnata Plugins
 License-File: LICENSE
 Author: James Weakley

{omnata_plugin_runtime-0.11.7a324 → omnata_plugin_runtime-0.11.7a325}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnata-plugin-runtime"
-version = "0.11.7-a324"
+version = "0.11.7-a325"
 description = "Classes and common runtime components for building and running Omnata Plugins"
 authors = ["James Weakley <james.weakley@omnata.com>"]
 readme = "README.md"

{omnata_plugin_runtime-0.11.7a324 → omnata_plugin_runtime-0.11.7a325}/src/omnata_plugin_runtime/omnata_plugin.py

@@ -49,6 +49,11 @@ from snowflake.snowpark.functions import col
 from tenacity import Retrying, stop_after_attempt, wait_fixed, retry_if_exception_message
 
 from .logging import OmnataPluginLogHandler, logger, tracer, meter
+stream_duration_gauge = meter.create_gauge(
+    name="omnata.sync_run.stream_duration",
+    description="The duration of stream processing",
+    unit="s",
+)
 from opentelemetry import context
 import math
 import numpy as np
@@ -1097,6 +1102,8 @@ class InboundSyncRequest(SyncRequest):
         self.state_register_table_name = results_table.get_fully_qualified_state_register_table_name()
         # this is keyed on stream name, each containing a list of dataframes and state updates mixed
         self._apply_results: Dict[str, List[pandas.DataFrame | Dict]] = {}
+        # track the start times of each stream, so we can calculate durations. The int is a epoch (time.time()) value
+        self._stream_start_times: Dict[str, int] = {}
 
     def apply_results_queue(self):
         """
@@ -1350,6 +1357,13 @@ class InboundSyncRequest(SyncRequest):
                 if sum([x.memory_usage(index=True).sum() for x in all_dfs if isinstance(x, pandas.DataFrame)]) > 200000000:
                     logger.debug(f"Applying criteria deletes queue immediately due to combined dataframe size")
                     self.apply_results_queue() 
+    
+    def mark_stream_started(self, stream_name: str):
+        """
+        Marks a stream as started, this is called automatically per stream when using @managed_inbound_processing.
+        """
+        logger.debug(f"Marking stream {stream_name} as started locally")
+        self._stream_start_times[stream_name] = time.time()
 
     def mark_stream_complete(self, stream_name: str):
         """
@@ -1357,6 +1371,20 @@ class InboundSyncRequest(SyncRequest):
         If @managed_inbound_processing is not used, call this whenever a stream has finished recieving records.
         """
         logger.debug(f"Marking stream {stream_name} as completed locally")
+        if stream_name in self._stream_start_times:
+            start_time = self._stream_start_times[stream_name]
+            duration = time.time() - start_time
+            stream_duration_gauge.set(
+                amount=duration,
+                attributes={
+                    "stream_name": stream_name,
+                    "sync_run_id": str(self._run_id),
+                    "sync_id": str(self._sync_id),
+                    "branch_name": str(self._branch_name) if self._branch_name is not None else 'main',
+                    "sync_direction": "inbound",
+                    "plugin_id": self.plugin_instance.get_manifest().plugin_id,
+                },
+            )
         with self._apply_results_lock:
             self._completed_streams.append(stream_name)
             # dedup just in case it's called twice
@@ -2333,12 +2361,8 @@ def __managed_inbound_processing_worker(
             sync_request: InboundSyncRequest = cast(
                 InboundSyncRequest, plugin_class_obj._sync_request
             )  # pylint: disable=protected-access
-            stream_duration_gauge = meter.create_gauge(
-                name="omnata.sync_run.stream_duration",
-                description="The duration of stream processing",
-                unit="s",
-            )
-            start_time = time.time()
+            if stream.stream_name not in sync_request._stream_start_times:
+                sync_request.mark_stream_started(stream.stream_name)
             # restore the first argument, was originally the dataframe/generator but now it's the appropriately sized dataframe
             try:
                 with tracer.start_as_current_span("managed_inbound_processing") as managed_inbound_processing_span:
@@ -2370,19 +2394,6 @@ def __managed_inbound_processing_worker(
                     omnata_plugin_logger.error(f"{type(e).__name__} syncing stream {stream.stream_name}", 
                                 exc_info=True, 
                                 extra={'stream_name':stream.stream_name})
-            finally:
-                duration = time.time() - start_time
-                stream_duration_gauge.set(
-                    amount=duration,
-                    attributes={
-                        "stream_name": stream.stream_name,
-                        "sync_run_id": str(sync_request._run_id),
-                        "sync_id": str(sync_request._sync_id),
-                        "branch_name": str(sync_request._branch_name) if sync_request._branch_name is not None else 'main',
-                        "sync_direction": "inbound",
-                        "plugin_id": plugin_class_obj.get_manifest().plugin_id,
-                    },
-                )
         except queue.Empty:
             logger.debug("streams queue is empty")
             return