omnata-plugin-runtime 0.8.0a188__tar.gz → 0.8.0a190__tar.gz

{omnata_plugin_runtime-0.8.0a188 → omnata_plugin_runtime-0.8.0a190}/PKG-INFO

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnata-plugin-runtime"
-version = "0.8.0-a188"
+version = "0.8.0-a190"
 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.8.0a188 → omnata_plugin_runtime-0.8.0a190}/src/omnata_plugin_runtime/configuration.py

@@ -19,7 +19,9 @@ if tuple(sys.version_info[:2]) >= (3, 9):
 else:
     # Python 3.8 and below
     from typing_extensions import Annotated
+from opentelemetry import trace
 
+tracer = trace.get_tracer(__name__)
 
 class MapperType(str, Enum):
     FIELD_MAPPING_SELECTOR = "field_mapping_selector"
@@ -871,7 +873,7 @@ InboundSyncStreamsConfiguration.model_rebuild()
 StoredFieldMappings.model_rebuild()
 OutboundSyncConfigurationParameters.model_rebuild()
 
-
+@tracer.start_as_current_span("get_secrets")
 def get_secrets(oauth_secret_name: Optional[str], other_secrets_name: Optional[str]
 ) -> Dict[str, StoredConfigurationValue]:
     connection_secrets = {}

{omnata_plugin_runtime-0.8.0a188 → omnata_plugin_runtime-0.8.0a190}/src/omnata_plugin_runtime/plugin_entrypoints.py

@@ -39,7 +39,7 @@ from .rate_limiting import ApiLimits, RateLimitState
 from opentelemetry import trace
 
 IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
-
+tracer = trace.get_tracer(__name__)
 
 class PluginEntrypoint:
     """
@@ -51,8 +51,7 @@ class PluginEntrypoint:
         self, plugin_fqn: str, session: Session, module_name: str, class_name: str
     ):
         logger.info(f"Initialising plugin entrypoint for {plugin_fqn}")
-        self.tracer = trace.get_tracer(__name__)
-        with self.tracer.start_as_current_span("plugin_initialization") as span:
+        with tracer.start_as_current_span("plugin_initialization") as span:
             self._session = session
             import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
             span.add_event("Adding plugin zip to path")
@@ -93,7 +92,7 @@ class PluginEntrypoint:
         logger.add_extra('omnata.sync_branch.id', request.sync_branch_id)
         logger.add_extra('omnata.sync_branch.name', request.sync_branch_name)
         logger.info("Entered sync method")
-        with self.tracer.start_as_current_span("sync_initialization") as span:
+        with tracer.start_as_current_span("initialization") as span:
             span.add_event("Fetching secrets")
 
             request = TypeAdapter(SyncRequestPayload).validate_python(sync_request)
@@ -186,11 +185,11 @@ class PluginEntrypoint:
             )
             try:
                 self._plugin_instance._configuration_parameters = parameters
-                with self.tracer.start_as_current_span("sync_execution") as span:
+                with tracer.start_as_current_span("invoke_plugin") as span:
                     with HttpRateLimiting(outbound_sync_request, parameters):
                         self._plugin_instance.sync_outbound(parameters, outbound_sync_request)
                 if self._plugin_instance.disable_background_workers is False:
-                    with self.tracer.start_as_current_span("results_finalization") as span:
+                    with tracer.start_as_current_span("results_finalization") as span:
                         outbound_sync_request.apply_results_queue()
                         outbound_sync_request.apply_rate_limit_state()
                 if outbound_sync_request.deadline_reached:
@@ -246,12 +245,12 @@ class PluginEntrypoint:
                 inbound_sync_request.update_activity("Invoking plugin")
                 logger.info(f"inbound sync request: {inbound_sync_request}")
                 # plugin_instance._inbound_sync_request = outbound_sync_request
-                with self.tracer.start_as_current_span("sync_execution") as span:
+                with tracer.start_as_current_span("invoke_plugin"):
                     with HttpRateLimiting(inbound_sync_request, parameters):
                         self._plugin_instance.sync_inbound(parameters, inbound_sync_request)
                 logger.info("Finished invoking plugin")
                 if self._plugin_instance.disable_background_workers is False:
-                    with self.tracer.start_as_current_span("results_finalization") as span:
+                    with tracer.start_as_current_span("results_finalization") as span:
                         inbound_sync_request.update_activity("Staging remaining records")
                         logger.info("Calling apply_results_queue")
                         inbound_sync_request.apply_results_queue()
@@ -304,6 +303,8 @@ class PluginEntrypoint:
         sync_parameters: Dict,
         current_form_parameters: Optional[Dict],
     ):
+        if function_name is None:
+            function_name = f"{sync_direction}_configuration_form"
         logger.add_extra('omnata.operation', 'configuration_form')
         logger.add_extra('omnata.connection.connectivity_option', connectivity_option)
         logger.add_extra('omnata.connection.connection_method', connection_method)
@@ -349,9 +350,10 @@ class PluginEntrypoint:
                 parameters.access_token_secret_name = oauth_secret_name
         the_function = getattr(
             self._plugin_instance,
-            function_name or f"{sync_direction}_configuration_form",
+            function_name
         )
-        script_result = the_function(parameters)
+        with tracer.start_as_current_span("invoke_plugin"):
+            script_result = the_function(parameters)
         if isinstance(script_result, BaseModel):
             script_result = script_result.model_dump()
         elif isinstance(script_result, List):
@@ -393,8 +395,8 @@ class PluginEntrypoint:
         )
         if oauth_secret_name is not None:
             parameters.access_token_secret_name = oauth_secret_name
-
-        script_result = self._plugin_instance.inbound_stream_list(parameters)
+        with tracer.start_as_current_span("invoke_plugin"):
+            script_result = self._plugin_instance.inbound_stream_list(parameters)
         if isinstance(script_result, BaseModel):
             script_result = script_result.model_dump()
         elif isinstance(script_result, List):
@@ -428,19 +430,21 @@ class PluginEntrypoint:
         logger.add_extra('omnata.connection.connectivity_option', connectivity_option)
         connectivity_option = TypeAdapter(ConnectivityOption).validate_python(connectivity_option)
         logger.info("Entered connection_form method")
-        if self._plugin_instance.connection_form.__code__.co_argcount==1:
-            form: List[ConnectionMethod] = self._plugin_instance.connection_form()
-        else:
-            form: List[ConnectionMethod] = self._plugin_instance.connection_form(connectivity_option)
+        with tracer.start_as_current_span("invoke_plugin"):
+            if self._plugin_instance.connection_form.__code__.co_argcount==1:
+                form: List[ConnectionMethod] = self._plugin_instance.connection_form()
+            else:
+                form: List[ConnectionMethod] = self._plugin_instance.connection_form(connectivity_option)
         return [f.model_dump() for f in form]
 
     def create_billing_events(self, session, event_request: Dict):
         logger.add_extra('omnata.operation', 'create_billing_events')
         logger.info("Entered create_billing_events method")
         request = TypeAdapter(BillingEventRequest).validate_python(event_request)
-        events: List[SnowflakeBillingEvent] = self._plugin_instance.create_billing_events(
-            request
-        )
+        with tracer.start_as_current_span("invoke_plugin"):
+            events: List[SnowflakeBillingEvent] = self._plugin_instance.create_billing_events(
+                request
+            )
         # create each billing event, waiting a second between each one
         first_time = True
         for billing_event in events:
@@ -527,31 +531,35 @@ class PluginEntrypoint:
         )
         if oauth_secret_name is not None:
             parameters.access_token_secret_name = oauth_secret_name
-        connect_response = self._plugin_instance.connect(
-            parameters=parameters
-        )
+        with tracer.start_as_current_span("invoke_plugin"):
+            connect_response = self._plugin_instance.connect(
+                parameters=parameters
+            )
         # the connect method can also return more network addresses. If so, we need to update the
         # network rule associated with the external access integration
         if connect_response is None:
             raise ValueError("Plugin did not return a ConnectResponse object from the connect method")
         if connect_response.network_addresses is not None:
-            existing_rule_result = self._session.sql(
-                f"desc network rule {network_rule_name}"
-            ).collect()
-            rule_values: List[str] = existing_rule_result[0].value_list.split(",")
-            rule_values = [r for r in rule_values if r != '']
-            logger.info(f"Existing rules for {network_rule_name}: {rule_values}")
-            for network_address in connect_response.network_addresses:
-                if network_address not in rule_values:
-                    rule_values.append(network_address)
-            #if len(rule_values)==0:
-            #    logger.info("No network addresses for plugin, adding localhost")
-            #    rule_values.append("https://localhost")
-            logger.info(f"New rules for {network_rule_name}: {rule_values}")
-            rule_values_string = ",".join([f"'{value}'" for value in rule_values])
-            self._session.sql(
-                f"alter network rule {network_rule_name} set value_list = ({rule_values_string})"
-            ).collect()
+            with tracer.start_as_current_span("network_rule_update") as network_rule_update_span:
+                network_rule_update_span.add_event("Retrieving existing network rule")
+                existing_rule_result = self._session.sql(
+                    f"desc network rule {network_rule_name}"
+                ).collect()
+                rule_values: List[str] = existing_rule_result[0].value_list.split(",")
+                rule_values = [r for r in rule_values if r != '']
+                logger.info(f"Existing rules for {network_rule_name}: {rule_values}")
+                for network_address in connect_response.network_addresses:
+                    if network_address not in rule_values:
+                        rule_values.append(network_address)
+                #if len(rule_values)==0:
+                #    logger.info("No network addresses for plugin, adding localhost")
+                #    rule_values.append("https://localhost")
+                logger.info(f"New rules for {network_rule_name}: {rule_values}")
+                rule_values_string = ",".join([f"'{value}'" for value in rule_values])
+                network_rule_update_span.add_event("Updating network rule")
+                self._session.sql(
+                    f"alter network rule {network_rule_name} set value_list = ({rule_values_string})"
+                ).collect()
 
         return connect_response.model_dump()
 
@@ -578,7 +586,8 @@ class PluginEntrypoint:
         )
         if oauth_secret_name is not None:
             connection_parameters.access_token_secret_name = oauth_secret_name
-        response: List[ApiLimits] = self._plugin_instance.api_limits(connection_parameters)
+        with tracer.start_as_current_span("invoke_plugin"):
+            response: List[ApiLimits] = self._plugin_instance.api_limits(connection_parameters)
         return [api_limit.model_dump() for api_limit in response]
 
     def outbound_record_validator(

{omnata_plugin_runtime-0.8.0a188 → omnata_plugin_runtime-0.8.0a190}/src/omnata_plugin_runtime/rate_limiting.py

@@ -19,6 +19,9 @@ from .logging import logger
 import pytz
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
+from opentelemetry import trace
+
+tracer = trace.get_tracer(__name__)
 
 TimeUnitType = Literal["second", "minute", "hour", "day"]
 
@@ -383,11 +386,12 @@ class RetryWithLogging(Retry):
         retry_after = self.get_retry_after(response)
         if retry_after:
             logger.info(f"Retrying after {retry_after} seconds due to Retry-After header")
-            if self.thread_cancellation_token is None:
-                time.sleep(retry_after)
-            else:
-                if self.thread_cancellation_token.wait(retry_after):
-                    raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting to expire")
+            with tracer.start_as_current_span("http_retry_wait"):
+                if self.thread_cancellation_token is None:
+                    time.sleep(retry_after)
+                else:
+                    if self.thread_cancellation_token.wait(retry_after):
+                        raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting to expire")
             return True
         return False
 
@@ -504,8 +508,9 @@ class RateLimitedSession(requests.Session):
                         raise InterruptedWhileWaitingException(message=f"The rate limiting wait time ({wait_time} seconds) would exceed the run deadline")
                     logger.info(f"Waiting for {wait_time} seconds before retrying {method} request to {url}")
                     # if wait() returns true, it means that the thread was cancelled
-                    if self.thread_cancellation_token.wait(wait_time):
-                        raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting to expire")
+                    with tracer.start_as_current_span("http_retry_wait"):
+                        if self.thread_cancellation_token.wait(wait_time):
+                            raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting to expire")
                 else:
                     current_url_retries = self.increment_retries(url)
                     if current_url_retries >= self.max_retries:
@@ -514,8 +519,9 @@ class RateLimitedSession(requests.Session):
                     if datetime.datetime.now(pytz.UTC) + datetime.timedelta(seconds=backoff_time) > self.run_deadline:
                         raise InterruptedWhileWaitingException(message=f"The rate limiting backoff time ({backoff_time} seconds) would exceed the run deadline")
                     logger.info(f"Waiting for {backoff_time} seconds before retrying {method} request to {url}")
-                    if self.thread_cancellation_token.wait(backoff_time):
-                        raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting backoff")
+                    with tracer.start_as_current_span("http_retry_wait"):
+                        if self.thread_cancellation_token.wait(backoff_time):
+                            raise InterruptedWhileWaitingException(message="The sync was interrupted while waiting for rate limiting backoff")
             else:
                 self.set_retries(url,0)  # Reset retries if the request is successful
                 return response