omnata-plugin-runtime 0.8.0a187__tar.gz → 0.8.0a189__tar.gz

{omnata_plugin_runtime-0.8.0a187 → omnata_plugin_runtime-0.8.0a189}/PKG-INFO

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

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

@@ -36,9 +36,10 @@ from .omnata_plugin import (
 )
 from pydantic import TypeAdapter
 from .rate_limiting import ApiLimits, RateLimitState
+from opentelemetry import trace
 
 IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
-
+tracer = trace.get_tracer(__name__)
 
 class PluginEntrypoint:
     """
@@ -50,46 +51,39 @@ class PluginEntrypoint:
         self, plugin_fqn: str, session: Session, module_name: str, class_name: str
     ):
         logger.info(f"Initialising plugin entrypoint for {plugin_fqn}")
-        self._session = session
-        import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
-        sys.path.append(os.path.join(import_dir, "app.zip"))
-        module = importlib.import_module(module_name)
-        class_obj = getattr(module, class_name)
-        self._plugin_instance: OmnataPlugin = class_obj()
-        self._plugin_instance._session = session  # pylint: disable=protected-access
-        # logging defaults
-        snowflake_logger = logging.getLogger("snowflake")
-        snowflake_logger.setLevel(logging.WARN) # we don't want snowflake queries being logged by default
-        # the sync engine can tell the plugin to override log level via a session variable
-        if session is not None:
-            try:
-                v = session.sql("select getvariable('LOG_LEVEL_OVERRIDES')").collect()
-                result = v[0][0]
-                if result is not None:
-                    log_level_overrides:Dict[str,str] = json.loads(result)
-                    for logger_name,level in log_level_overrides.items():
-                        logger_override = logging.getLogger(logger_name)
-                        logger_override.setLevel(level)
-                        logger_override.propagate = False
-                        for handler in logger_override.handlers:
-                            handler.setLevel(level)
-            except Exception as e:
-                logger.error(f"Error setting log level overrides: {str(e)}")
+        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")
+            sys.path.append(os.path.join(import_dir, "app.zip"))
+            span.add_event("Importing plugin module")
+            module = importlib.import_module(module_name)
+            class_obj = getattr(module, class_name)
+            self._plugin_instance: OmnataPlugin = class_obj()
+            self._plugin_instance._session = session  # pylint: disable=protected-access
+            # logging defaults
+            snowflake_logger = logging.getLogger("snowflake")
+            snowflake_logger.setLevel(logging.WARN) # we don't want snowflake queries being logged by default
+            # the sync engine can tell the plugin to override log level via a session variable
+            if session is not None:
+                try:
+                    span.add_event("Checking log level overrides")
+                    v = session.sql("select getvariable('LOG_LEVEL_OVERRIDES')").collect()
+                    result = v[0][0]
+                    if result is not None:
+                        log_level_overrides:Dict[str,str] = json.loads(result)
+                        span.add_event("Applying log level overrides",log_level_overrides)
+                        for logger_name,level in log_level_overrides.items():
+                            logger_override = logging.getLogger(logger_name)
+                            logger_override.setLevel(level)
+                            logger_override.propagate = False
+                            for handler in logger_override.handlers:
+                                handler.setLevel(level)
+                except Exception as e:
+                    logger.error(f"Error setting log level overrides: {str(e)}")
         
 
     def sync(self, sync_request: Dict):
-        logger.info("Entered sync method")
-        request = TypeAdapter(SyncRequestPayload).validate_python(sync_request)
-        connection_secrets = get_secrets(
-            request.oauth_secret_name, request.other_secrets_name
-        )
-        omnata_log_handler = OmnataPluginLogHandler(
-            session=self._session,
-            sync_id=request.sync_id,
-            sync_branch_id=request.sync_branch_id,
-            connection_id=request.connection_id,
-            sync_run_id=request.run_id,
-        )
         logger.add_extra('omnata.operation', 'sync')
         logger.add_extra('omnata.sync.id', request.sync_id)
         logger.add_extra('omnata.sync.direction', request.sync_direction)
@@ -97,50 +91,67 @@ class PluginEntrypoint:
         logger.add_extra('omnata.sync.run_id', request.run_id)
         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 tracer.start_as_current_span("initialization") as span:
+            span.add_event("Fetching secrets")
 
-        omnata_log_handler.register(
-            request.logging_level, self._plugin_instance.additional_loggers()
-        )
-        # construct some connection parameters for the purpose of getting the api limits
-        connection_parameters = ConnectionConfigurationParameters(
-            connection_method=request.connection_method,
-            connectivity_option=request.connectivity_option,
-            connection_parameters=request.connection_parameters,
-            connection_secrets=connection_secrets
-        )
-        if request.oauth_secret_name is not None:
-            connection_parameters.access_token_secret_name = request.oauth_secret_name
-        all_api_limits = self._plugin_instance.api_limits(connection_parameters)
-        logger.info(
-            f"Default API limits: {json.dumps(to_jsonable_python(all_api_limits))}"
-        )
-        all_api_limits_by_category = {
-            api_limit.endpoint_category: api_limit for api_limit in all_api_limits
-        }
-        all_api_limits_by_category.update(
-            {
-                k: v
-                for k, v in [
-                    (x.endpoint_category, x) for x in request.api_limit_overrides
-                ]
+            request = TypeAdapter(SyncRequestPayload).validate_python(sync_request)
+            connection_secrets = get_secrets(
+                request.oauth_secret_name, request.other_secrets_name
+            )
+            span.add_event("Configuring log handler")
+            omnata_log_handler = OmnataPluginLogHandler(
+                session=self._session,
+                sync_id=request.sync_id,
+                sync_branch_id=request.sync_branch_id,
+                connection_id=request.connection_id,
+                sync_run_id=request.run_id,
+            )
+
+            omnata_log_handler.register(
+                request.logging_level, self._plugin_instance.additional_loggers()
+            )
+            # construct some connection parameters for the purpose of getting the api limits
+            connection_parameters = ConnectionConfigurationParameters(
+                connection_method=request.connection_method,
+                connectivity_option=request.connectivity_option,
+                connection_parameters=request.connection_parameters,
+                connection_secrets=connection_secrets
+            )
+            if request.oauth_secret_name is not None:
+                connection_parameters.access_token_secret_name = request.oauth_secret_name
+            span.add_event("Configuring API Limits")
+            all_api_limits = self._plugin_instance.api_limits(connection_parameters)
+            logger.info(
+                f"Default API limits: {json.dumps(to_jsonable_python(all_api_limits))}"
+            )
+            all_api_limits_by_category = {
+                api_limit.endpoint_category: api_limit for api_limit in all_api_limits
             }
-        )
-        api_limits = list(all_api_limits_by_category.values())
-        return_dict = {}
-        logger.info(
-            f"Rate limits state: {json.dumps(to_jsonable_python(request.rate_limits_state))}"
-        )
-        (rate_limit_state_all, rate_limit_state_this_branch) = RateLimitState.collapse(request.rate_limits_state,request.sync_id, request.sync_branch_name)
-        # if any endpoint categories have no state, give them an empty state
-        for api_limit in api_limits:
-            if api_limit.endpoint_category not in rate_limit_state_all:
-                rate_limit_state_all[api_limit.endpoint_category] = RateLimitState(
-                    wait_until=None, previous_request_timestamps=[]
-                )
-            if api_limit.endpoint_category not in rate_limit_state_this_branch:
-                rate_limit_state_this_branch[api_limit.endpoint_category] = RateLimitState(
-                    wait_until=None, previous_request_timestamps=[]
-                )
+            all_api_limits_by_category.update(
+                {
+                    k: v
+                    for k, v in [
+                        (x.endpoint_category, x) for x in request.api_limit_overrides
+                    ]
+                }
+            )
+            api_limits = list(all_api_limits_by_category.values())
+            return_dict = {}
+            logger.info(
+                f"Rate limits state: {json.dumps(to_jsonable_python(request.rate_limits_state))}"
+            )
+            (rate_limit_state_all, rate_limit_state_this_branch) = RateLimitState.collapse(request.rate_limits_state,request.sync_id, request.sync_branch_name)
+            # if any endpoint categories have no state, give them an empty state
+            for api_limit in api_limits:
+                if api_limit.endpoint_category not in rate_limit_state_all:
+                    rate_limit_state_all[api_limit.endpoint_category] = RateLimitState(
+                        wait_until=None, previous_request_timestamps=[]
+                    )
+                if api_limit.endpoint_category not in rate_limit_state_this_branch:
+                    rate_limit_state_this_branch[api_limit.endpoint_category] = RateLimitState(
+                        wait_until=None, previous_request_timestamps=[]
+                    )
 
         if request.sync_direction == "outbound":
             parameters = OutboundSyncConfigurationParameters(
@@ -174,11 +185,13 @@ class PluginEntrypoint:
             )
             try:
                 self._plugin_instance._configuration_parameters = parameters
-                with HttpRateLimiting(outbound_sync_request, parameters):
-                    self._plugin_instance.sync_outbound(parameters, outbound_sync_request)
+                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:
-                    outbound_sync_request.apply_results_queue()
-                    outbound_sync_request.apply_rate_limit_state()
+                    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:
                     # 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
@@ -232,19 +245,21 @@ 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 HttpRateLimiting(inbound_sync_request, parameters):
-                    self._plugin_instance.sync_inbound(parameters, inbound_sync_request)
+                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:
-                    inbound_sync_request.update_activity("Staging remaining records")
-                    logger.info("Calling apply_results_queue")
-                    inbound_sync_request.apply_results_queue()
-                    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)}")
+                    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()
+                        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
@@ -288,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)
@@ -333,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):
@@ -377,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):
@@ -412,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:
@@ -511,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()
 
@@ -562,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(