omnata-plugin-runtime 0.8.0a195__tar.gz → 0.8.0a198__tar.gz

{omnata_plugin_runtime-0.8.0a195 → omnata_plugin_runtime-0.8.0a198}/PKG-INFO

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

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

@@ -48,6 +48,7 @@ from snowflake.snowpark.functions import col
 from tenacity import Retrying, stop_after_attempt, wait_fixed, retry_if_exception_message
 
 from .logging import OmnataPluginLogHandler, logger, tracer
+from opentelemetry import context
 
 from .api import (
     PluginMessage,
@@ -316,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:
@@ -387,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()
@@ -399,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()
@@ -408,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.
@@ -1637,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
@@ -1946,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()}"
     )
@@ -2128,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(

{omnata_plugin_runtime-0.8.0a195 → omnata_plugin_runtime-0.8.0a198}/src/omnata_plugin_runtime/rate_limiting.py

@@ -15,13 +15,10 @@ 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
+from .logging import logger, tracer
 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"]