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

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

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

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

@@ -11,6 +11,7 @@ from enum import Enum
 
 from abc import ABC
 from pydantic import BaseModel, Field, PrivateAttr, SerializationInfo, TypeAdapter, field_validator, model_serializer, validator  # pylint: disable=no-name-in-module
+from .logging import logger
 
 if tuple(sys.version_info[:2]) >= (3, 9):
     # Python 3.9 and above
@@ -19,8 +20,6 @@ else:
     # Python 3.8 and below
     from typing_extensions import Annotated
 
-logger = logging.getLogger(__name__)
-
 
 class MapperType(str, Enum):
     FIELD_MAPPING_SELECTOR = "field_mapping_selector"

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

@@ -24,7 +24,7 @@ from .configuration import (
     ConnectivityOption
 )
 from .forms import ConnectionMethod, FormInputField, FormOption
-from .logging import OmnataPluginLogHandler
+from .logging import OmnataPluginLogHandler, logger
 from .omnata_plugin import (
     SnowflakeBillingEvent,
     BillingEventRequest,
@@ -36,9 +36,7 @@ from .omnata_plugin import (
 )
 from pydantic import TypeAdapter
 from .rate_limiting import ApiLimits, RateLimitState
-
-# set the logger class to our custom logger so that pydantic errors are handled correctly
-logger = logging.getLogger(__name__)
+from opentelemetry import trace
 
 IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
 
@@ -53,89 +51,108 @@ 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)}")
+        self.tracer = trace.get_tracer(__name__)
+        with self.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.add_extra('omnata.operation', 'sync')
+        logger.add_extra('omnata.sync.id', request.sync_id)
+        logger.add_extra('omnata.sync.direction', request.sync_direction)
+        logger.add_extra('omnata.connection.id', request.connection_id)
+        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")
-        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,
-        )
-        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
-                ]
+        with self.tracer.start_as_current_span("sync_initialization") as span:
+            span.add_event("Fetching secrets")
+
+            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(
@@ -169,11 +186,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 self.tracer.start_as_current_span("sync_execution") 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 self.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
@@ -227,19 +246,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 self.tracer.start_as_current_span("sync_execution") as span:
+                    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 self.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
@@ -283,6 +304,12 @@ class PluginEntrypoint:
         sync_parameters: Dict,
         current_form_parameters: Optional[Dict],
     ):
+        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)
+        logger.add_extra('omnata.configuration_form.function_name', function_name)
+        logger.add_extra('omnata.sync.direction', sync_direction)
+
         logger.info("Entered configuration_form method")
         sync_strategy = normalise_nulls(sync_strategy)
         oauth_secret_name = normalise_nulls(oauth_secret_name)
@@ -342,6 +369,10 @@ class PluginEntrypoint:
         sync_parameters: Dict,
         selected_streams: Optional[List[str]], # None to return all streams without requiring schema
     ):
+        logger.add_extra('omnata.operation', 'list_streams')
+        logger.add_extra('omnata.connection.connectivity_option', connectivity_option)
+        logger.add_extra('omnata.connection.connection_method', connection_method)
+        logger.add_extra('omnata.sync.direction', 'inbound')
         logger.debug("Entered list_streams method")
         oauth_secret_name = normalise_nulls(oauth_secret_name)
         other_secrets_name = normalise_nulls(other_secrets_name)
@@ -393,6 +424,8 @@ class PluginEntrypoint:
         return results
 
     def connection_form(self,connectivity_option: str):
+        logger.add_extra('omnata.operation', 'connection_form')
+        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:
@@ -402,6 +435,7 @@ class PluginEntrypoint:
         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(
@@ -474,6 +508,9 @@ class PluginEntrypoint:
         oauth_secret_name: Optional[str],
         other_secrets_name: Optional[str],
     ):
+        logger.add_extra('omnata.operation', 'connection_test')
+        logger.add_extra('omnata.connection.connectivity_option', connectivity_option)
+        logger.add_extra('omnata.connection.connection_method', method)
         logger.info("Entered connect method")
         logger.info(f"Connection parameters: {connection_parameters}")
         connection_secrets = get_secrets(oauth_secret_name, other_secrets_name)
@@ -524,6 +561,9 @@ class PluginEntrypoint:
                     connection_parameters: Dict,
                     oauth_secret_name: Optional[str],
                     other_secrets_name: Optional[str]):
+        logger.add_extra('omnata.operation', 'api_limits')
+        logger.add_extra('omnata.connection.connectivity_option', connectivity_option)
+        logger.add_extra('omnata.connection.connection_method', method)
         logger.info("Entered api_limits method")
         connection_secrets = get_secrets(oauth_secret_name, other_secrets_name)
         from omnata_plugin_runtime.omnata_plugin import (

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

@@ -15,12 +15,11 @@ import logging
 from pydantic import Field, root_validator, PrivateAttr, field_serializer
 from pydantic_core import to_jsonable_python
 from .configuration import SubscriptableBaseModel
+from .logging import logger
 import pytz
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
-logger = getLogger(__name__)
-
 TimeUnitType = Literal["second", "minute", "hour", "day"]
 
 HttpMethodType = Literal[