omnata-plugin-runtime 0.3.28a85__tar.gz → 0.4.0__tar.gz

{omnata_plugin_runtime-0.3.28a85 → omnata_plugin_runtime-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omnata-plugin-runtime
-Version: 0.3.28a85
+Version: 0.4.0
 Summary: Classes and common runtime components for building and running Omnata Plugins
 Author: James Weakley
 Author-email: james.weakley@omnata.com

{omnata_plugin_runtime-0.3.28a85 → omnata_plugin_runtime-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnata-plugin-runtime"
-version = "0.3.28-a85"
+version = "0.4.0"
 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.3.28a85 → omnata_plugin_runtime-0.4.0}/src/omnata_plugin_runtime/api.py

@@ -66,8 +66,9 @@ class PluginMessageStreamProgressUpdate(BaseModel):
     message_type: Literal["stream_record_counts"] = "stream_record_counts"
     stream_total_counts: Dict[str, int]
     completed_streams: List[str]
-    # older runtime versions didn't have this, so the sync engine can't expect it
+    # older runtime versions didn't have these, so the sync engine can't expect it
     stream_errors: Optional[Dict[str,str]]
+    total_records_estimate: Optional[Dict[str,int]]
 
 
 class PluginMessageCancelledStreams(BaseModel):

{omnata_plugin_runtime-0.3.28a85 → omnata_plugin_runtime-0.4.0}/src/omnata_plugin_runtime/omnata_plugin.py

@@ -906,6 +906,7 @@ class InboundSyncRequest(SyncRequest):
         self._stream_record_counts: Dict[str, int] = {
             stream_name: 0 for stream_name in self._streams_dict.keys()
         }
+        self._total_records_estimate: Optional[Dict[str,int]] = {}
         self._stream_change_counts: Dict[str, int] = {
             stream_name: 0 for stream_name in self._streams_dict.keys()
         }
@@ -933,6 +934,7 @@ class InboundSyncRequest(SyncRequest):
                     ] = self._stream_record_counts[stream_name] + total_length
                     results.extend(non_empty_dfs)  # remove any None/empty dataframes
                     stream_names.append(stream_name)
+                self._session.sql("begin").collect()
                 if len(results) > 0:
                     logger.info(
                         f"Applying {len(results)} batches of queued results"
@@ -943,6 +945,7 @@ class InboundSyncRequest(SyncRequest):
                     self._apply_results_dataframe(stream_names, all_dfs)
                 # update the stream state object too
                 self._apply_latest_states()
+                self._session.sql("commit").collect()
                 for stream_name in stream_names:
                     self._apply_results[stream_name] = None
                 self._apply_results = {}
@@ -975,7 +978,8 @@ class InboundSyncRequest(SyncRequest):
         new_progress_update = PluginMessageStreamProgressUpdate(
                 stream_total_counts=self._stream_record_counts,
                 completed_streams=self._completed_streams,
-                stream_errors=self._omnata_log_handler.stream_global_errors
+                stream_errors=self._omnata_log_handler.stream_global_errors,
+                total_records_estimate=self._total_records_estimate
             )
         if self._last_stream_progress_update is None or new_progress_update != self._last_stream_progress_update:
             self._plugin_message(
@@ -1109,6 +1113,12 @@ class InboundSyncRequest(SyncRequest):
         self._completed_streams.append(stream_name)
         # dedup just in case it's called twice
         self._completed_streams = list(set(self._completed_streams))
+    
+    def set_stream_record_count(self, stream_name: str, count: int):
+        """
+        Sets the record count for a stream, used to provide progress updates.
+        """
+        self._stream_record_counts[stream_name] = count
 
     def _enqueue_state(self, stream_name: str, new_state: Any):
         """

{omnata_plugin_runtime-0.3.28a85 → omnata_plugin_runtime-0.4.0}/src/omnata_plugin_runtime/rate_limiting.py

@@ -132,19 +132,15 @@ class ApiLimits(SubscriptableBaseModel):
         ):
             longest_wait = rate_limit_state.wait_until
         for request_rate in self.request_rates:
-            print(request_rate)
             if rate_limit_state.previous_request_timestamps is not None and len(rate_limit_state.previous_request_timestamps) > 0:
                 previous_request_timestamps = rate_limit_state.get_relevant_history(request_rate)
-                print(f"previous_request_timestamps: {len(previous_request_timestamps)}")
                 request_index = request_rate.request_count - 1
                 if request_index > len(previous_request_timestamps) - 1:
                     continue # we have not yet made enough requests to hit this rate limit
                 request_index = len(previous_request_timestamps) - 1
                 timestamp_at_horizon = previous_request_timestamps[request_index]
-                print("timestamp_at_horizon: " + str(timestamp_at_horizon))
                 now = datetime.datetime.now().astimezone(datetime.timezone.utc)
                 seconds_since_horizon = (timestamp_at_horizon - now).total_seconds()
-                print("seconds_since_horizon: " + str(seconds_since_horizon))
                 next_allowed_request = timestamp_at_horizon + datetime.timedelta(
                     seconds=request_rate.number_of_seconds()
                 )