omnata-plugin-runtime 0.7.0a185__tar.gz → 0.8.0__tar.gz

{omnata_plugin_runtime-0.7.0a185 → omnata_plugin_runtime-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omnata-plugin-runtime
-Version: 0.7.0a185
+Version: 0.8.0
 Summary: Classes and common runtime components for building and running Omnata Plugins
 Author: James Weakley
 Author-email: james.weakley@omnata.com
@@ -19,9 +19,11 @@ Requires-Dist: idna (<=3.7)
 Requires-Dist: jinja2 (>=3.1.2,<=3.1.4)
 Requires-Dist: markupsafe (<=2.1.3)
 Requires-Dist: numpy (<1.27.0)
+Requires-Dist: opentelemetry-api (<=1.23.0)
 Requires-Dist: packaging (<=24.1)
 Requires-Dist: pandas (<=2.2.2)
 Requires-Dist: platformdirs (<=3.10.0)
+Requires-Dist: protobuf (<=4.25.3)
 Requires-Dist: pyarrow (<=16.1.0)
 Requires-Dist: pycparser (<=2.21)
 Requires-Dist: pydantic (>=2,<=2.8.2)
@@ -34,10 +36,12 @@ Requires-Dist: requests (>=2,<=2.32.3)
 Requires-Dist: setuptools (<=72.1.0)
 Requires-Dist: snowflake-connector-python (>=3,<=3.12.0)
 Requires-Dist: snowflake-snowpark-python (==1.23.0)
+Requires-Dist: snowflake-telemetry-python (<=0.5.0)
 Requires-Dist: tenacity (>=8,<=8.2.3)
 Requires-Dist: tomlkit (<=0.11.1)
 Requires-Dist: urllib3 (<=2.2.2)
 Requires-Dist: wheel (<=0.43.0)
+Requires-Dist: wrapt (<=1.14.1)
 Description-Content-Type: text/markdown
 
 # omnata-plugin-runtime

{omnata_plugin_runtime-0.7.0a185 → omnata_plugin_runtime-0.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnata-plugin-runtime"
-version = "0.7.0-a185"
+version = "0.8.0"
 description = "Classes and common runtime components for building and running Omnata Plugins"
 authors = ["James Weakley <james.weakley@omnata.com>"]
 readme = "README.md"
@@ -38,6 +38,10 @@ wheel = "<=0.43.0" # latest version available on Snowflake Anaconda
 pyyaml = "<=6.0.1" # latest version available on Snowflake Anaconda
 cffi = "<=1.16.0" # latest version available on Snowflake Anaconda
 pyarrow = "<=16.1.0" # latest version available on Snowflake Anaconda
+wrapt = "<=1.14.1" # latest version available on Snowflake Anaconda
+opentelemetry-api = "<=1.23.0" # latest version available on Snowflake Anaconda
+snowflake-telemetry-python = "<=0.5.0" # latest version available on Snowflake Anaconda
+protobuf = "<=4.25.3" # latest version available on Snowflake Anaconda
 
 [tool.poetry.dev-dependencies]
 pytest = "^6.2.4"

{omnata_plugin_runtime-0.7.0a185 → omnata_plugin_runtime-0.8.0}/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, tracer
 
 if tuple(sys.version_info[:2]) >= (3, 9):
     # Python 3.9 and above
@@ -19,9 +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"
     JINJA_TEMPLATE = "jinja_template"
@@ -872,7 +870,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.7.0a185 → omnata_plugin_runtime-0.8.0}/src/omnata_plugin_runtime/logging.py

@@ -4,11 +4,39 @@ Custom logging functionality for Omnata
 import logging
 import logging.handlers
 import traceback
-from logging import Logger
+from logging import Logger, getLogger
 from typing import Dict, List, Optional
 from snowflake.snowpark import Session
 from pydantic import ValidationError
+from snowflake import telemetry
+from opentelemetry import trace
 
+tracer = trace.get_tracer('omnata_plugin_runtime')
+
+class CustomLoggerAdapter(logging.LoggerAdapter):
+    """
+    A logger adapter which attaches current trace and span IDs to log messages, so that they can be correlated to traces.
+    Also offers the ability to add static extras.
+    """
+    def __init__(self, logger:logging.Logger, extra):
+        super(CustomLoggerAdapter, self).__init__(logger, extra)
+        self.extra_extras = {}
+    
+    def add_extra(self, key, value):
+        self.extra_extras[key] = value
+
+    def process(self, msg, kwargs):
+        extra = kwargs.get("extra", {})
+        current_span = trace.get_current_span()
+        context = current_span.get_span_context() if current_span is not None else None
+        if context is not None:
+            extra["trace_id"] = format(context.trace_id, 'x') # format as hex string to be consistent with Snowflake's handler
+            extra["span_id"] = format(context.span_id, 'x')
+        extra.update(self.extra_extras)
+        kwargs["extra"] = extra
+        return msg, kwargs
+
+logger = CustomLoggerAdapter(getLogger("omnata_plugin"), {})
 
 def log_exception(exception: Exception, logger_instance: Logger):
     """

{omnata_plugin_runtime-0.7.0a185 → omnata_plugin_runtime-0.8.0}/src/omnata_plugin_runtime/omnata_plugin.py

@@ -47,7 +47,8 @@ from snowflake.snowpark import Session
 from snowflake.snowpark.functions import col
 from tenacity import Retrying, stop_after_attempt, wait_fixed, retry_if_exception_message
 
-from .logging import OmnataPluginLogHandler
+from .logging import OmnataPluginLogHandler, logger, tracer
+from opentelemetry import context
 
 from .api import (
     PluginMessage,
@@ -88,7 +89,6 @@ from .rate_limiting import (
     RateLimitedSession
 )
 
-logger = getLogger(__name__)
 SortDirectionType = Literal["asc", "desc"]
 
 
@@ -317,6 +317,8 @@ class SyncRequest(ABC):
         self._cancel_checking_task = None
         self._rate_limit_update_task = None
         self._last_stream_progress_update = None
+        # store the opentelemetry context so that it can be attached inside threads
+        self.opentelemetry_context = context.get_current()
 
         threading.excepthook = self.thread_exception_hook
         if self.development_mode is False:
@@ -388,6 +390,7 @@ class SyncRequest(ABC):
         """
         Designed to be run in a thread, this method polls the results every 20 seconds and sends them back to Snowflake.
         """
+        context.attach(self.opentelemetry_context)
         while not cancellation_token.is_set():
             logger.debug("apply results worker checking for results")
             self.apply_results_queue()
@@ -400,6 +403,7 @@ class SyncRequest(ABC):
         It also gives us the latest rate limit state from Snowflake, so that activity on other syncs/branches can
         impact rate limiting on this one.
         """
+        context.attach(self.opentelemetry_context)
         while not cancellation_token.is_set():
             try:
                 self.apply_rate_limit_state()
@@ -409,6 +413,7 @@ class SyncRequest(ABC):
         logger.info("rate limit update worker exiting")
 
     def __cancel_checking_worker(self, cancellation_token:threading.Event):
+        context.attach(self.opentelemetry_context)
         """
         Designed to be run in a thread, this method checks to see if the sync run has been cancelled
         or if the deadline has been reached.
@@ -810,22 +815,23 @@ class OutboundSyncRequest(SyncRequest):
         logger.debug("applying results to table")
         # use a random table name with a random string to avoid collisions
         with self._snowflake_query_lock:
-            for attempt in Retrying(stop=stop_after_attempt(30),wait=wait_fixed(2),reraise=True,retry=retry_if_exception_message(match=".*(is being|was) committed.*")):
-                with attempt:
-                    success, nchunks, nrows, _ = write_pandas(
-                        conn=self._session._conn._cursor.connection,  # pylint: disable=protected-access
-                        df=self._preprocess_results_dataframe(results_df),
-                        quote_identifiers=False,
-                        table_name=self._full_results_table_name,
-                        auto_create_table=False
-                    )
-                    if not success:
-                        raise ValueError(
-                            f"Failed to write results to table {self._full_results_table_name}"
+            with tracer.start_as_current_span("apply_results"):
+                for attempt in Retrying(stop=stop_after_attempt(30),wait=wait_fixed(2),reraise=True,retry=retry_if_exception_message(match=".*(is being|was) committed.*")):
+                    with attempt:
+                        success, nchunks, nrows, _ = write_pandas(
+                            conn=self._session._conn._cursor.connection,  # pylint: disable=protected-access
+                            df=self._preprocess_results_dataframe(results_df),
+                            quote_identifiers=False,
+                            table_name=self._full_results_table_name,
+                            auto_create_table=False
+                        )
+                        if not success:
+                            raise ValueError(
+                                f"Failed to write results to table {self._full_results_table_name}"
+                            )
+                        logger.debug(
+                            f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
                         )
-                    logger.debug(
-                        f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
-                    )
 
     def __dataframe_wrapper(
         self, data_frame: pandas.DataFrame, render_jinja: bool = True
@@ -1465,36 +1471,37 @@ class InboundSyncRequest(SyncRequest):
         """
         if len(results_df) > 0:
             with self._snowflake_query_lock:
-                for attempt in Retrying(stop=stop_after_attempt(30),wait=wait_fixed(2),reraise=True,retry=retry_if_exception_message(match=".*(is being|was) committed.*")):
-                    with attempt:
-                        logger.debug(
-                            f"Applying {len(results_df)} results to {self._full_results_table_name}"
-                        )
-                        # try setting parquet engine here, since the engine parameter does not seem to make it through to the write_pandas function
-                        success, nchunks, nrows, _ = write_pandas(
-                            conn=self._session._conn._cursor.connection,  # pylint: disable=protected-access
-                            df=results_df,
-                            table_name=self._full_results_table_name,
-                            quote_identifiers=False,  # already done in get_temp_table_name
-                            # schema='INBOUND_RAW', # it seems to be ok to provide schema in the table name
-                            table_type="transient"
-                        )
-                        if not success:
-                            raise ValueError(
-                                f"Failed to write results to table {self._full_results_table_name}"
+                with tracer.start_as_current_span("apply_results"):
+                    for attempt in Retrying(stop=stop_after_attempt(30),wait=wait_fixed(2),reraise=True,retry=retry_if_exception_message(match=".*(is being|was) committed.*")):
+                        with attempt:
+                            logger.debug(
+                                f"Applying {len(results_df)} results to {self._full_results_table_name}"
                             )
-                        logger.debug(
-                            f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
-                        )
-                        # temp tables aren't allowed
-                        # snowflake_df = self._session.create_dataframe(results_df)
-                        # snowflake_df.write.save_as_table(table_name=temp_table,
-                        #                                mode='append',
-                        #                                column_order='index',
-                        #                                #create_temp_table=True
-                        #                                )
-                        for stream_name in stream_names:
-                            self._results_exist[stream_name] = True
+                            # try setting parquet engine here, since the engine parameter does not seem to make it through to the write_pandas function
+                            success, nchunks, nrows, _ = write_pandas(
+                                conn=self._session._conn._cursor.connection,  # pylint: disable=protected-access
+                                df=results_df,
+                                table_name=self._full_results_table_name,
+                                quote_identifiers=False,  # already done in get_temp_table_name
+                                # schema='INBOUND_RAW', # it seems to be ok to provide schema in the table name
+                                table_type="transient"
+                            )
+                            if not success:
+                                raise ValueError(
+                                    f"Failed to write results to table {self._full_results_table_name}"
+                                )
+                            logger.debug(
+                                f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
+                            )
+                            # temp tables aren't allowed
+                            # snowflake_df = self._session.create_dataframe(results_df)
+                            # snowflake_df.write.save_as_table(table_name=temp_table,
+                            #                                mode='append',
+                            #                                column_order='index',
+                            #                                #create_temp_table=True
+                            #                                )
+                            for stream_name in stream_names:
+                                self._results_exist[stream_name] = True
         else:
             logger.debug("Results dataframe is empty, not applying")
 
@@ -1636,6 +1643,8 @@ class OmnataPlugin(ABC):
         take care of various status updates and loading enqueued results.
         Only set this to True if you plan to do all of this yourself (e.g. in a Java stored proc)
         """
+        # store the opentelemetry context so that it can be attached inside threads
+        self.opentelemetry_context = context.get_current()
     
 
     @abstractmethod
@@ -1945,6 +1954,7 @@ def __managed_outbound_processing_worker(
     Consumes a fixed sized set of records by passing them to the wrapped function,
     while adhering to the defined API constraints.
     """
+    context.attach(plugin_class_obj.opentelemetry_context)
     logger.debug(
         f"worker {worker_index} processing. Cancelled: {cancellation_token.is_set()}"
     )
@@ -2127,6 +2137,7 @@ def __managed_inbound_processing_worker(
     A worker thread for the managed_inbound_processing annotation.
     Passes single streams at a time to the wrapped function, adhering to concurrency constraints.
     """
+    context.attach(plugin_class_obj.opentelemetry_context)
     while not cancellation_token.is_set():
         # Get our generator object out of the queue
         logger.debug(
@@ -2137,14 +2148,16 @@ def __managed_inbound_processing_worker(
             logger.debug(f"stream returned from queue: {stream}")
             # restore the first argument, was originally the dataframe/generator but now it's the appropriately sized dataframe
             try:
-                logger.debug(f"worker {worker_index} processing stream {stream.stream_name}, invoking plugin class method {method.__name__}")
-                result = method(plugin_class_obj, *(stream, *method_args), **method_kwargs)
-                logger.debug(f"worker {worker_index} completed processing stream {stream.stream_name}")
-                if result is not None and result is False:
-                    logger.info(f"worker {worker_index} requested that {stream.stream_name} be not marked as complete")
-                else:
-                    logger.info(f"worker {worker_index} marking stream {stream.stream_name} as complete")
-                    plugin_class_obj._sync_request.mark_stream_complete(stream.stream_name)
+                with tracer.start_as_current_span("managed_inbound_processing") as managed_inbound_processing_span:
+                    logger.debug(f"worker {worker_index} processing stream {stream.stream_name}, invoking plugin class method {method.__name__}")
+                    managed_inbound_processing_span.set_attribute("stream_name", stream.stream_name)
+                    result = method(plugin_class_obj, *(stream, *method_args), **method_kwargs)
+                    logger.debug(f"worker {worker_index} completed processing stream {stream.stream_name}")
+                    if result is not None and result is False:
+                        logger.info(f"worker {worker_index} requested that {stream.stream_name} be not marked as complete")
+                    else:
+                        logger.info(f"worker {worker_index} marking stream {stream.stream_name} as complete")
+                        plugin_class_obj._sync_request.mark_stream_complete(stream.stream_name)
             except InterruptedWhileWaitingException:
                 # If an inbound run is cancelled while waiting for rate limiting, this should mean that
                 # the cancellation is handled elsewhere, so we don't need to do anything special here other than stop waiting

{omnata_plugin_runtime-0.7.0a185 → omnata_plugin_runtime-0.8.0}/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, tracer
 from .omnata_plugin import (
     SnowflakeBillingEvent,
     BillingEventRequest,
@@ -36,13 +36,10 @@ 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"
 
-
 class PluginEntrypoint:
     """
     This class gives each plugin's stored procs an initial point of contact.
@@ -53,89 +50,107 @@ 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,
-        )
-        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
-                ]
+        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")
+        with tracer.start_as_current_span("initialization") as span:
+            span.add_event("Fetching secrets")
+
+            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 +184,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
@@ -227,19 +244,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
@@ -283,6 +302,14 @@ 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)
+        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)
@@ -322,9 +349,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):
@@ -342,6 +370,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)
@@ -362,8 +394,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):
@@ -393,20 +425,25 @@ 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:
-            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:
@@ -474,6 +511,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)
@@ -490,31 +530,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()
 
@@ -524,6 +568,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 (
@@ -538,7 +585,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.7.0a185 → omnata_plugin_runtime-0.8.0}/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, tracer
 import pytz
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 
-logger = getLogger(__name__)
-
 TimeUnitType = Literal["second", "minute", "hour", "day"]
 
 HttpMethodType = Literal[
@@ -384,11 +383,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
 
@@ -505,8 +505,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:
@@ -515,8 +516,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