omnata-plugin-runtime 0.4.4__tar.gz → 0.4.5__tar.gz

{omnata_plugin_runtime-0.4.4 → omnata_plugin_runtime-0.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omnata-plugin-runtime
-Version: 0.4.4
+Version: 0.4.5
 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.4.4 → omnata_plugin_runtime-0.4.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnata-plugin-runtime"
-version = "0.4.4"
+version = "0.4.5"
 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.4.4 → omnata_plugin_runtime-0.4.5}/src/omnata_plugin_runtime/plugin_entrypoints.py

@@ -12,7 +12,7 @@ from pydantic import BaseModel, parse_obj_as  # pylint: disable=no-name-in-modul
 import pydantic.json
 from snowflake.snowpark import Session
 
-from .api import SyncRequestPayload, handle_proc_result
+from .api import PluginMessageStreamProgressUpdate, SyncRequestPayload, handle_proc_result
 from .configuration import (
     ConnectionConfigurationParameters,
     InboundSyncConfigurationParameters,
@@ -218,11 +218,22 @@ class PluginEntrypoint:
                     inbound_sync_request.update_activity("Staging remaining records")
                     logger.info("Calling apply_results_queue")
                     inbound_sync_request.apply_results_queue()
-                    logger.info("Calling apply_rate_limit_state")
-                    inbound_sync_request.apply_rate_limit_state()
-                    logger.info("Calling apply_progress_updates")
-                    # we can't ignore errors here, because the sync engine needs to get the progress updates before we return
-                    inbound_sync_request.apply_progress_updates(ignore_errors=False)
+                    try:
+                        # this is not critical, we wouldn't fail the sync over rate limit usage capture
+                        logger.info("Calling apply_rate_limit_state")
+                        inbound_sync_request.apply_rate_limit_state()
+                    except Exception as e:
+                        logger.error(f"Error applying rate limit state: {str(e)}")
+                    # here we used to do a final inbound_sync_request.apply_progress_updates(ignore_errors=False)
+                    # but it was erroring too much since there was usually a lot of DDL activity on the Snowflake side
+                    # so instead, we'll provide a final progress update via a return value from the proc
+                final_progress_update = PluginMessageStreamProgressUpdate(
+                    stream_total_counts=inbound_sync_request._stream_record_counts,
+                    completed_streams=inbound_sync_request._completed_streams,
+                    stream_errors=omnata_log_handler.stream_global_errors,
+                    total_records_estimate=inbound_sync_request._total_records_estimate
+                )
+                return_dict["final_progress_update"] = final_progress_update.dict()
                 if inbound_sync_request.deadline_reached:
                     # if we actually hit the deadline, this is flagged by the cancellation checking worker and the cancellation
                     # token is set. We throw it here as an error since that's currently how it flows back to the engine with a DELAYED state
@@ -239,11 +250,6 @@ class PluginEntrypoint:
                         inbound_sync_request._rate_limit_update_task.join()  # pylint: disable=protected-access
                 except Exception as e:
                     logger.error(f"Error cleaning up threading: {str(e)}")
-            return_dict["errored_streams"] = list(
-                omnata_log_handler.stream_global_errors.keys()
-            )
-            # we need to calculate counts for:
-            # CHANGED_COUNT by counting up the records in INBOUND_STREAM_RECORD_COUNTS
         logger.info("Finished applying records")
         return return_dict