omnata-plugin-runtime 0.11.6__py3-none-any.whl → 0.12.2a347__py3-none-any.whl

omnata_plugin_runtime/configuration.py

@@ -693,6 +693,9 @@ class ConnectionConfigurationParameters(SubscriptableBaseModel):
     _snowflake: Optional[Any] = PrivateAttr(  # or use Any to annotate the type and use Field to initialize
         default=None
     )
+    _plugin_instance: Optional[Any] = PrivateAttr(  # Reference to OmnataPlugin instance for accessing sync_request
+        default=None
+    )
 
     @model_validator(mode='after')
     def validate_ngrok_tunnel_settings(self) -> Self:
@@ -739,6 +742,22 @@ class ConnectionConfigurationParameters(SubscriptableBaseModel):
         """
         if parameter_name=='access_token' and self.access_token_secret_name is not None:
             import _snowflake # pylint: disable=import-error, import-outside-toplevel # type: ignore
+            from .threading_utils import is_managed_worker_thread
+            
+            # Check if we're in a worker thread using the explicit flag
+            # This is more reliable than checking thread names
+            if is_managed_worker_thread() and self._plugin_instance is not None and self._plugin_instance._sync_request is not None:
+                logger.debug(f"Worker thread requesting access_token via OAuth token service for secret: {self.access_token_secret_name}")
+                try:
+                    access_token = self._plugin_instance._sync_request.request_access_token_from_main_thread(
+                        self.access_token_secret_name
+                    )
+                    return StoredConfigurationValue(value=access_token)
+                except Exception as e:
+                    logger.error(f"Error requesting access_token from main thread: {e}")
+                    raise
+            
+            # Otherwise, call _snowflake directly (main thread)
             return StoredConfigurationValue(
                 value=_snowflake.get_oauth_access_token(self.access_token_secret_name)
             )
@@ -1005,14 +1024,41 @@ 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]
+def get_secrets(oauth_secret_name: Optional[str], other_secrets_name: Optional[str],
+                sync_request: Optional[Any] = None
 ) -> Dict[str, StoredConfigurationValue]:
+    """
+    Get secrets from Snowflake. This function can be called from the main thread or worker threads.
+    When called from worker threads (e.g., within @managed_inbound_processing) for OAuth access tokens,
+    it will automatically route the OAuth token request through the main thread to avoid threading issues
+    with _snowflake.get_oauth_access_token. Other secrets can be fetched directly.
+    
+    :param oauth_secret_name: The name of the OAuth secret to retrieve
+    :param other_secrets_name: The name of other secrets to retrieve  
+    :param sync_request: Optional SyncRequest instance for worker threads. If not provided, will attempt to detect.
+    :return: Dictionary of StoredConfigurationValue objects
+    """
+    from .threading_utils import is_managed_worker_thread
     connection_secrets = {}
     import _snowflake # pylint: disable=import-error, import-outside-toplevel # type: ignore
+    
+    # OAuth token needs special handling in worker threads
     if oauth_secret_name is not None:
-        connection_secrets["access_token"] = StoredConfigurationValue(
-            value=_snowflake.get_oauth_access_token(oauth_secret_name)
-        )
+        if is_managed_worker_thread() and sync_request is not None:
+            logger.debug(f"Worker thread requesting OAuth access token via main thread for secret: {oauth_secret_name}")
+            try:
+                access_token = sync_request.request_access_token_from_main_thread(oauth_secret_name)
+                connection_secrets["access_token"] = StoredConfigurationValue(value=access_token)
+            except Exception as e:
+                logger.error(f"Error requesting OAuth access token from main thread: {e}")
+                raise
+        else:
+            # Main thread - call _snowflake directly
+            connection_secrets["access_token"] = StoredConfigurationValue(
+                value=_snowflake.get_oauth_access_token(oauth_secret_name)
+            )
+    
+    # Other secrets can be fetched directly from any thread
     if other_secrets_name is not None:
         try:
             secret_string_content = _snowflake.get_generic_secret_string(

omnata_plugin_runtime/json_schema.py

@@ -886,8 +886,9 @@ def prune(view_part: SnowflakeViewPart, joined_parts: List[SnowflakeViewPart]) -
                             deps.append(dep_key)
                         else:
                             logger.warning(
-                                f"Column {column.original_name} in {part.stream_name} references "
-                                f"{ref_field} in {resolved_stream}, which doesn't exist"
+                                msg=f"Column {column.original_name} in {part.stream_name} references "
+                                f"{ref_field} in {resolved_stream}, which doesn't exist",
+                                stack_info=True
                             )
                             has_invalid_dep = True
             
@@ -906,8 +907,9 @@ def prune(view_part: SnowflakeViewPart, joined_parts: List[SnowflakeViewPart]) -
                 for dep_key in deps:
                     if dep_key in columns_with_invalid_deps:
                         logger.warning(
-                            f"Column {col_key[1]} in {col_key[0]} depends on "
-                            f"{dep_key[1]} in {dep_key[0]}, which has invalid dependencies"
+                            msg=f"Column {col_key[1]} in {col_key[0]} depends on "
+                            f"{dep_key[1]} in {dep_key[0]}, which has invalid dependencies",
+                            stack_info=True
                         )
                         columns_with_invalid_deps.add(col_key)
                         changed = True
@@ -995,8 +997,9 @@ def prune(view_part: SnowflakeViewPart, joined_parts: List[SnowflakeViewPart]) -
                 
                 if missing_refs:
                     logger.warning(
-                        f"Removing column {col.original_name} from {part.stream_name} because it references "
-                        f"non-existent column(s): {', '.join(missing_refs)}"
+                        msg=f"Removing column {col.original_name} from {part.stream_name} because it references "
+                        f"non-existent column(s): {', '.join(missing_refs)}",
+                        stack_info=True
                     )
                 else:
                     # Column is not needed (not referenced by main part)

omnata_plugin_runtime/omnata_plugin.py

@@ -15,7 +15,7 @@ if tuple(sys.version_info[:2]) >= (3, 9):
 else:
     # Python 3.8 and below
     from typing_extensions import Annotated
-
+from dataclasses import dataclass
 import zipfile
 import datetime
 import http
@@ -49,6 +49,11 @@ 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, meter
+stream_duration_gauge = meter.create_gauge(
+    name="omnata.sync_run.stream_duration",
+    description="The duration of stream processing",
+    unit="s",
+)
 from opentelemetry import context
 import math
 import numpy as np
@@ -96,6 +101,7 @@ from .rate_limiting import (
 from .json_schema import (
     FullyQualifiedTable
 )
+from .threading_utils import is_managed_worker_thread, set_managed_worker_thread
 
 SortDirectionType = Literal["asc", "desc"]
 
@@ -265,6 +271,29 @@ def jinja_filter(func):
     func.is_jinja_filter = True
     return func
 
+@dataclass
+class StateResult:
+    """
+    Represents the current cursor state of a stream. This simple wrapper just helps us identify what type of
+    object is in the apply_results list.
+    """
+    new_state: Any
+
+@dataclass
+class RecordsToUploadResult:
+    """
+    Represents the records to upload for a stream. This simple wrapper just helps us identify what type of
+    object is in the apply_results list.
+    """
+    records: pandas.DataFrame
+
+@dataclass
+class CriteriaDeleteResult:
+    """
+    Represents the result of processing criteria deletes for a stream. This simple wrapper just helps us identify what type of
+    object is in the apply_results list.
+    """
+    criteria_deletes: pandas.DataFrame
 
 class SyncRequest(ABC):
     """
@@ -347,6 +376,11 @@ class SyncRequest(ABC):
         self._last_states_update = None
         # store the opentelemetry context so that it can be attached inside threads
         self.opentelemetry_context = context.get_current()
+        
+        # Secrets service for thread-safe access to _snowflake.get_oauth_access_token
+        # which can only be called from the main thread
+        # The main thread (in decorator wait loops) will service these requests
+        self._secrets_request_queue: queue.Queue = queue.Queue()
 
         threading.excepthook = self.thread_exception_hook
         if self.development_mode is False:
@@ -471,6 +505,82 @@ class SyncRequest(ABC):
             cancellation_token.wait(20)
         logger.info("cancel checking worker exiting")
 
+    def _service_oauth_token_request(self):
+        """
+        Services any pending OAuth token requests from worker threads.
+        This should be called periodically from the main thread while waiting for workers.
+        Returns True if any requests were serviced, False otherwise.
+        """
+        import _snowflake  # pylint: disable=import-error, import-outside-toplevel # type: ignore
+        
+        serviced_any = False
+        # Process all pending requests (non-blocking)
+        while not self._secrets_request_queue.empty():
+            try:
+                request = self._secrets_request_queue.get_nowait()
+            except queue.Empty:
+                break
+                
+            serviced_any = True
+            oauth_secret_name = request.get('oauth_secret_name')
+            response_queue = request['response_queue']
+            
+            logger.debug(f"Main thread servicing OAuth token request for secret: {oauth_secret_name}")
+            
+            try:
+                # Call _snowflake.get_oauth_access_token directly (we're on the main thread now)
+                access_token = _snowflake.get_oauth_access_token(oauth_secret_name)
+                
+                # Send successful response
+                response_queue.put({
+                    'success': True,
+                    'result': access_token
+                })
+            except Exception as e:
+                logger.error(f"Error servicing OAuth token request: {e}")
+                # Send error response
+                response_queue.put({
+                    'success': False,
+                    'error': str(e)
+                })
+            finally:
+                self._secrets_request_queue.task_done()
+        
+        return serviced_any
+
+    def request_access_token_from_main_thread(self, oauth_secret_name: str, timeout: int = 30) -> str:
+        """
+        Request OAuth access token from the main thread. This should be called from worker threads
+        when they need to access the OAuth token via _snowflake.get_oauth_access_token.
+        The main thread services these requests while waiting for workers to complete.
+        
+        :param oauth_secret_name: The name of the OAuth secret to retrieve
+        :param timeout: Maximum time to wait for the response in seconds
+        :return: The OAuth access token string
+        :raises TimeoutError: if the request times out
+        :raises ValueError: if the secrets service returns an error
+        """
+        # Create a response queue for this specific request
+        response_queue: queue.Queue = queue.Queue()
+        
+        logger.debug(f"Requesting OAuth access token from main thread for secret: {oauth_secret_name}")
+        
+        # Put the request in the queue with its own response queue
+        self._secrets_request_queue.put({
+            'oauth_secret_name': oauth_secret_name,
+            'response_queue': response_queue
+        })
+        
+        # Block on the response queue with timeout
+        try:
+            response = response_queue.get(timeout=timeout)
+            if response['success']:
+                return response['result']
+            else:
+                raise ValueError(f"Error getting OAuth access token: {response['error']}")
+        except queue.Empty:
+            raise TimeoutError(f"Timeout waiting for OAuth access token request after {timeout} seconds")
+
     @abstractmethod
     def apply_results_queue(self):
         """
@@ -1057,7 +1167,6 @@ class InboundSyncRequest(SyncRequest):
         }
         
         # These are similar to the results, but represent requests to delete records by some criteria
-        self._apply_results_criteria_deletes: Dict[str, List[pandas.DataFrame]] = {}
         self._temp_tables = {}
         self._temp_table_lock = threading.Lock()
         self._results_exist: Dict[
@@ -1096,7 +1205,9 @@ class InboundSyncRequest(SyncRequest):
         self._criteria_deletes_table_name = results_table.get_fully_qualified_criteria_deletes_table_name()
         self.state_register_table_name = results_table.get_fully_qualified_state_register_table_name()
         # this is keyed on stream name, each containing a list of dataframes and state updates mixed
-        self._apply_results: Dict[str, List[pandas.DataFrame | Dict]] = {}
+        self._apply_results: Dict[str, List[RecordsToUploadResult | StateResult | CriteriaDeleteResult]] = {}
+        # track the start times of each stream, so we can calculate durations. The int is a epoch (time.time()) value
+        self._stream_start_times: Dict[str, int] = {}
 
     def apply_results_queue(self):
         """
@@ -1105,7 +1216,8 @@ class InboundSyncRequest(SyncRequest):
         logger.debug("InboundSyncRequest apply_results_queue")
         if self._apply_results is not None:
             with self._apply_results_lock:
-                results:List[pandas.DataFrame] = []
+                records_to_upload:List[pandas.DataFrame] = []
+                criteria_deletes_to_upload:List[pandas.DataFrame] = []
                 stream_states_for_upload:Dict[str, Dict[str, Any]] = {}
                 for stream_name, stream_results in self._apply_results.items():
                     # the stream results contains an ordered sequence of dataframes and state updates (append only)
@@ -1113,9 +1225,9 @@ class InboundSyncRequest(SyncRequest):
                     # so first, we iterate backwards to find the last state update
                     last_state_index = -1
                     for i in range(len(stream_results) - 1, -1, -1):
-                        if isinstance(stream_results[i], dict):
+                        if isinstance(stream_results[i], StateResult):
                             last_state_index = i
-                            stream_states_for_upload[stream_name] = stream_results[i]
+                            stream_states_for_upload[stream_name] = stream_results[i].new_state
                             break
                     # if there are no state updates, we can't do anything with this stream
                     if last_state_index == -1:
@@ -1124,56 +1236,54 @@ class InboundSyncRequest(SyncRequest):
                         )
                         continue
                     assert isinstance(stream_states_for_upload[stream_name], dict), "Latest state must be a dictionary"
-                    # now we can take the dataframes up to the last state update
-                    dfs = stream_results[:last_state_index]
-                    non_empty_dfs = [
-                        x for x in dfs if x is not None and isinstance(x, pandas.DataFrame) and len(x) > 0
+                    # now we can take the record dataframes up to the last state update
+                    results_subset = stream_results[:last_state_index]
+                    non_empty_record_dfs:List[pandas.DataFrame] = [
+                        x.records for x in results_subset 
+                        if x is not None and isinstance(x, RecordsToUploadResult) and len(x.records) > 0
                     ]
                     # get the total length of all the dataframes
-                    total_length = sum([len(x) for x in non_empty_dfs])
+                    total_length = sum([len(x) for x in non_empty_record_dfs])
                     # add the count of this batch to the total for this stream
                     self._stream_record_counts[
                         stream_name
                     ] = self._stream_record_counts[stream_name] + total_length
-                    results.extend(non_empty_dfs)
+                    records_to_upload.extend(non_empty_record_dfs)
+                    # also handle any criteria deletes
+                    criteria_deletes_to_upload.extend([
+                        x.criteria_deletes for x in results_subset
+                        if x is not None and isinstance(x, CriteriaDeleteResult) and len(x.criteria_deletes) > 0
+                    ])
                     # now remove everything up to the last state update
                     # we do this so that we don't apply the same state update multiple times
+                    # keep everything after the last state update
                     self._apply_results[stream_name] = stream_results[
                         last_state_index + 1 :
-                    ]  # keep everything after the last state update
-                if len(results) > 0:
-                    logger.debug(
-                        f"Applying {len(results)} batches of queued results"
-                    )
-                    # upload all cached apply results
-                    all_dfs = pandas.concat(results)
-                    query_id = self._apply_results_dataframe(list(stream_states_for_upload.keys()), all_dfs)
-                    # now that the results have been updated, we need to insert records into the state register table
-                    # we do this by inserting the latest state for each stream
+                    ]
+                
+                if len(records_to_upload) > 0 or len(criteria_deletes_to_upload) > 0:
+                    if len(records_to_upload) > 0:
+                        logger.debug(
+                            f"Applying {len(records_to_upload)} batches of queued results"
+                        )
+                        # upload all cached apply results
+                        records_to_upload_combined = pandas.concat(records_to_upload)
+                        self._apply_results_dataframe(list(stream_states_for_upload.keys()), records_to_upload_combined)
+                        # now that the results have been updated, we need to insert records into the state register table
+                        # we do this by inserting the latest state for each stream
+                    if len(criteria_deletes_to_upload) > 0:
+                        logger.debug(
+                            f"Applying {len(criteria_deletes_to_upload)} batches of queued criteria deletes"
+                        )
+                        # upload all cached apply results
+                        all_criteria_deletes = pandas.concat(criteria_deletes_to_upload)
+                        self._apply_criteria_deletes_dataframe(all_criteria_deletes)
+                    
+                    query_id = self._get_query_id_for_now()
                     self._directly_insert_to_state_register(
                         stream_states_for_upload, query_id=query_id
                     )
         
-        # also take care of uploading delete requests
-        # technically these should be managed along with the state, however there aren't any scenarios where checkpointing is done
-        # and deletes have an impact. This is because we only checkpoint in scenarios where the target table is empty first
-        if hasattr(self,'_apply_results_criteria_deletes') and self._apply_results_criteria_deletes is not None:
-            with self._apply_results_lock:
-                results:List[pandas.DataFrame] = []
-                for stream_name, stream_results in self._apply_results_criteria_deletes.items():
-                    results.extend([
-                        x for x in stream_results if x is not None and len(x) > 0
-                    ])
-                if len(results) > 0:
-                    logger.debug(
-                        f"Applying {len(results)} batches of queued criteria deletes"
-                    )
-                    # upload all cached apply results
-                    all_dfs = pandas.concat(results)
-                    self._apply_criteria_deletes_dataframe(all_dfs)
-                    # clear the delete requests
-                    self._apply_results_criteria_deletes = {}
-        
 
         # update the inbound stream record counts, so we can see progress
         # we do this last, because marking a stream as completed will cause the sync engine to process it
@@ -1281,42 +1391,58 @@ class InboundSyncRequest(SyncRequest):
         if stream_name is None or len(stream_name) == 0:
             raise ValueError("Stream name cannot be empty")
         with self._apply_results_lock:
-            existing_results: List[pandas.DataFrame] = []
+            existing_results: List[RecordsToUploadResult | StateResult | CriteriaDeleteResult] = []
             if stream_name in self._apply_results:
                 existing_results = self._apply_results[stream_name]
-            existing_results.append(self._preprocess_results_list(stream_name, results, is_delete))
+            existing_results.append(RecordsToUploadResult(
+                records=self._preprocess_results_list(stream_name, results, is_delete)
+            ))
             if new_state is not None:
-                existing_results.append(new_state)  # append the new state at the end
+                existing_results.append(
+                    StateResult(new_state=new_state)
+                )  # append the new state at the end
             self._apply_results[stream_name] = existing_results
-        # if the total size of all the dataframes exceeds 200MB, apply the results immediately
-        # we'll use df.memory_usage(index=True) for this
         if self.development_mode is False:
             # note: we want to do it for all values in self._apply_results, not just the new one
-            # so first we need to get the list of lists from the dictionary values and flatten it
-            # then we can sum the memory usage of each dataframe
-            # if the total exceeds 200MB, we apply the results immediately
-            all_df_lists:List[List[pandas.DataFrame]] = list(self._apply_results.values())
-            # flatten
-            all_dfs:List[pandas.DataFrame] = [x for sublist in all_df_lists for x in sublist if isinstance(x, pandas.DataFrame)]
-            combined_length = sum([len(x) for x in all_dfs])
-            # first, don't bother if the count is less than 10000, since it's unlikely to be even close
-            if combined_length > 10000:
-                if sum([x.memory_usage(index=True).sum() for x in all_dfs]) > 200000000:
-                    logger.debug(f"Applying results queue immediately due to combined dataframe size")
-                    self.apply_results_queue()
+            self._apply_results_if_size_exceeded()
+    
+    def _apply_results_if_size_exceeded(self,):
+        # so first we need to get the list of lists from the dictionary values and flatten it
+        # then we can sum the memory usage of each dataframe
+        # if the total exceeds 200MB, we apply the results immediately
+        all_df_lists:List[List[RecordsToUploadResult | StateResult | CriteriaDeleteResult]] = list(self._apply_results.values())
+        # flatten
+        all_dfs:List[pandas.DataFrame] = []
+        for sublist in all_df_lists:
+            for x in sublist:
+                if isinstance(x, RecordsToUploadResult):
+                    all_dfs.append(x.records)
+                if isinstance(x, CriteriaDeleteResult):
+                    all_dfs.append(x.criteria_deletes)
+        combined_length = sum([len(x) for x in all_dfs])
+        # first, don't bother if the count is less than 10000, since it's unlikely to be even close
+        if combined_length > 10000:
+            if sum([x.memory_usage(index=True).sum() for x in all_dfs]) > 200000000:
+                logger.debug(f"Applying results queue immediately due to combined dataframe size")
+                self.apply_results_queue()
     
-    def delete_by_criteria(self, stream_name: str, criteria: Dict[str, Any]):
+    def delete_by_criteria(self, stream_name: str, criteria: Dict[str, Any], new_state: Any):
         """
-        Submits some critera (field→value dict) which will cause matching records to be marked as deleted at the end of the run. 
+        Submits some critera (field→value dict) which will cause matching records to be marked as deleted
+         during checkpointing or at the end of the run. 
         This feature was created primarily for array fields that become child streams.
         The parent record is updated, which means there is a set of new children, but we need to delete the previously sync'd records and we don't know their identifiers.
 
-        The criteria is applied before the new records for the current run are applied. In other words, it will not delete any records from the current run.
+        The criteria is applied before the new records for the current run/checkpoint are applied.
 
         For a record to be deleted, it must match fields with all the criteria supplied. At least one field value must be provided.
+
+        If you pass in None for new_state, then the criteria delete will not apply unless you also enqueue record state for the same stream. This provides the ability to do an atomic delete-and-replace.
+        If you pass in some new state, then the criteria deletes will be applied in isolation along with the new state in a transaction.
         """
         if len(criteria) == 0:
             raise ValueError("At least one field value must be provided for deletion criteria")
+        
         if stream_name not in self._streams_dict:
             raise ValueError(
                 f"Cannot delete records for stream {stream_name} as its configuration doesn't exist"
@@ -1329,27 +1455,26 @@ class InboundSyncRequest(SyncRequest):
             logger.debug(
                 f"Enqueuing {len(criteria)} delete criteria for stream {stream_name} for upload"
             )
-            existing_results: List[pandas.DataFrame] = []
-            if stream_name in self._apply_results_criteria_deletes:
-                existing_results = self._apply_results_criteria_deletes[stream_name]
-            existing_results.append(pandas.DataFrame([{"STREAM_NAME":stream_name,"DELETE_CRITERIA": criteria}]))
-            self._apply_results_criteria_deletes[stream_name] = existing_results
-        # if the total size of all the dataframes exceeds 200MB, apply the results immediately
-        # we'll use df.memory_usage(index=True) for this
+            existing_results: List[RecordsToUploadResult | StateResult | CriteriaDeleteResult] = []
+            if stream_name in self._apply_results:
+                existing_results = self._apply_results[stream_name]
+            existing_results.append(
+                CriteriaDeleteResult(
+                    criteria_deletes=pandas.DataFrame([{"STREAM_NAME":stream_name,"DELETE_CRITERIA": criteria}])))
+            if new_state is not None:
+                existing_results.append(
+                    StateResult(new_state=new_state)
+                )  # append the new state at the end
+            self._apply_results[stream_name] = existing_results
         if self.development_mode is False:
-            # note: we want to do it for all values in self._apply_results_criteria_deletes, not just the new one
-            # so first we need to get the list of lists from the dictionary values and flatten it
-            # then we can sum the memory usage of each dataframe
-            # if the total exceeds 200MB, we apply the results immediately
-            all_df_lists:List[List[pandas.DataFrame]] = list(self._apply_results_criteria_deletes.values())
-            # flatten
-            all_dfs:List[pandas.DataFrame] = [x for sublist in all_df_lists for x in sublist]
-            combined_length = sum([len(x) for x in all_dfs])
-            # first, don't both if the count is less than 10000, since it's unlikely to be even close
-            if combined_length > 10000:
-                if sum([x.memory_usage(index=True).sum() for x in all_dfs if isinstance(x, pandas.DataFrame)]) > 200000000:
-                    logger.debug(f"Applying criteria deletes queue immediately due to combined dataframe size")
-                    self.apply_results_queue() 
+            self._apply_results_if_size_exceeded()
+    
+    def mark_stream_started(self, stream_name: str):
+        """
+        Marks a stream as started, this is called automatically per stream when using @managed_inbound_processing.
+        """
+        logger.debug(f"Marking stream {stream_name} as started locally")
+        self._stream_start_times[stream_name] = time.time()
 
     def mark_stream_complete(self, stream_name: str):
         """
@@ -1357,6 +1482,20 @@ class InboundSyncRequest(SyncRequest):
         If @managed_inbound_processing is not used, call this whenever a stream has finished recieving records.
         """
         logger.debug(f"Marking stream {stream_name} as completed locally")
+        if stream_name in self._stream_start_times:
+            start_time = self._stream_start_times[stream_name]
+            duration = time.time() - start_time
+            stream_duration_gauge.set(
+                amount=duration,
+                attributes={
+                    "stream_name": stream_name,
+                    "sync_run_id": str(self._run_id),
+                    "sync_id": str(self._sync_id),
+                    "branch_name": str(self._branch_name) if self._branch_name is not None else 'main',
+                    "sync_direction": "inbound",
+                    "plugin_id": self.plugin_instance.get_manifest().plugin_id,
+                },
+            )
         with self._apply_results_lock:
             self._completed_streams.append(stream_name)
             # dedup just in case it's called twice
@@ -1402,7 +1541,7 @@ class InboundSyncRequest(SyncRequest):
         with self._apply_results_lock:
             if stream_name in self._apply_results:
                 if len(self._apply_results[stream_name]) > 0:
-                    self._apply_results[stream_name].append(new_state)
+                    self._apply_results[stream_name].append(StateResult(new_state=new_state))
                     return
         
         self._directly_insert_to_state_register(
@@ -1463,7 +1602,7 @@ class InboundSyncRequest(SyncRequest):
             logger.debug(f"Failure to convert inbound data: {str(exception)}")
         return data
 
-    def _preprocess_results_list(self, stream_name: str, results: List[Dict],is_delete:Union[bool,List[bool]]):
+    def _preprocess_results_list(self, stream_name: str, results: List[Dict],is_delete:Union[bool,List[bool]]) -> pandas.DataFrame:
         """
         Creates a dataframe from the enqueued list, ready to upload.
         The result is a dataframe contain all (and only):
@@ -1608,7 +1747,7 @@ class InboundSyncRequest(SyncRequest):
         hash_object = hashlib.sha256(key_string.encode())
         return hash_object.hexdigest()
 
-    def _apply_results_dataframe(self, stream_names: List[str], results_df: pandas.DataFrame) -> Optional[str]:
+    def _apply_results_dataframe(self, stream_names: List[str], results_df: pandas.DataFrame):
         """
         Applies results for an inbound sync. The results are staged into a temporary
         table in Snowflake, so that we can make an atomic commit at the end.
@@ -1635,7 +1774,6 @@ class InboundSyncRequest(SyncRequest):
                                 raise ValueError(
                                     f"Failed to write results to table {self._full_results_table_name}"
                                 )
-                            query_id = self._get_query_id_for_now()
                             logger.debug(
                                 f"Wrote {nrows} rows and {nchunks} chunks to table {self._full_results_table_name}"
                             )
@@ -1648,7 +1786,6 @@ class InboundSyncRequest(SyncRequest):
                             #                                )
                             for stream_name in stream_names:
                                 self._results_exist[stream_name] = True
-                            return query_id
         else:
             logger.debug("Results dataframe is empty, not applying")
 
@@ -2138,6 +2275,9 @@ 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.
     """
+    # Mark this thread as a managed worker thread
+    set_managed_worker_thread(True)
+    
     context.attach(plugin_class_obj.opentelemetry_context)
     logger.debug(
         f"worker {worker_index} processing. Cancelled: {cancellation_token.is_set()}"
@@ -2277,6 +2417,8 @@ def managed_outbound_processing(concurrency: int, batch_size: int):
                         task.join()  # Ensure the thread is fully finished
                         tasks.remove(task)
                         logger.info(f"Thread {task.name} has completed processing")
+                # Service any OAuth token requests from worker threads while we wait
+                self._sync_request._service_oauth_token_request()
                 time.sleep(1)  # Avoid busy waiting
             logger.info("All workers completed processing")
 
@@ -2321,6 +2463,9 @@ 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.
     """
+    # Mark this thread as a managed worker thread
+    set_managed_worker_thread(True)
+    
     context.attach(plugin_class_obj.opentelemetry_context)
     while not cancellation_token.is_set():
         # Get our generator object out of the queue
@@ -2333,12 +2478,8 @@ def __managed_inbound_processing_worker(
             sync_request: InboundSyncRequest = cast(
                 InboundSyncRequest, plugin_class_obj._sync_request
             )  # pylint: disable=protected-access
-            stream_duration_counter = meter.create_histogram(
-                name="omnata.sync_run.stream_duration",
-                description="The duration of stream processing",
-                unit="s",
-            )
-            start_time = time.time()
+            if stream.stream_name not in sync_request._stream_start_times:
+                sync_request.mark_stream_started(stream.stream_name)
             # restore the first argument, was originally the dataframe/generator but now it's the appropriately sized dataframe
             try:
                 with tracer.start_as_current_span("managed_inbound_processing") as managed_inbound_processing_span:
@@ -2370,19 +2511,6 @@ def __managed_inbound_processing_worker(
                     omnata_plugin_logger.error(f"{type(e).__name__} syncing stream {stream.stream_name}", 
                                 exc_info=True, 
                                 extra={'stream_name':stream.stream_name})
-            finally:
-                duration = time.time() - start_time
-                stream_duration_counter.record(
-                    duration,
-                    attributes={
-                        "stream_name": stream.stream_name,
-                        "sync_run_id": str(sync_request._run_id),
-                        "sync_id": str(sync_request._sync_id),
-                        "branch_name": str(sync_request._branch_name) if sync_request._branch_name is not None else 'main',
-                        "sync_direction": "inbound",
-                        "plugin_id": plugin_class_obj.get_manifest().plugin_id,
-                    },
-                )
         except queue.Empty:
             logger.debug("streams queue is empty")
             return
@@ -2476,6 +2604,8 @@ def managed_inbound_processing(concurrency: int):
                             task.join()  # Ensure the thread is fully finished
                             tasks.remove(task)
                             logger.info(f"Thread {task.name} has completed processing")
+                    # Service any OAuth token requests from worker threads while we wait
+                    self._sync_request._service_oauth_token_request()
                     time.sleep(1)  # Avoid busy waiting
                 logger.info("All workers completed processing")
 
@@ -2733,7 +2863,7 @@ def omnata_udf(
     
     return decorator
 
-def find_udf_functions(path:str = '.',top_level_modules:Optional[List[str]] = None) -> List[UDFDefinition]:
+def find_udf_functions(path:str = '.',top_level_modules:Optional[List[str]] = None, exclude_top_level_modules:Optional[List[str]] = None) -> List[UDFDefinition]:
     """
     Finds all functions in the specified directory which have the 'omnata_udf' decorator applied
     """
@@ -2749,6 +2879,9 @@ def find_udf_functions(path:str = '.',top_level_modules:Optional[List[str]] = No
         if top_level_modules is not None:
             if len([x for x in top_level_modules if module_name.startswith(x)]) == 0:
                 continue
+        if exclude_top_level_modules is not None:
+            if any(module_name.startswith(y) for y in exclude_top_level_modules):
+                continue
         module = importlib.import_module(module_name)
 
         # Iterate over all members of the module

omnata_plugin_runtime/plugin_entrypoints.py

@@ -188,6 +188,8 @@ class PluginEntrypoint:
                 sync_id=request.sync_id,
                 branch_name=request.sync_branch_name
             )
+            # Store plugin_instance reference in parameters for worker thread OAuth token access
+            parameters._plugin_instance = self._plugin_instance  # pylint: disable=protected-access
             try:
                 self._plugin_instance._configuration_parameters = parameters
                 with tracer.start_as_current_span("invoke_plugin") as span:
@@ -246,6 +248,8 @@ class PluginEntrypoint:
                 sync_id=request.sync_id,
                 branch_name=request.sync_branch_name
             )
+            # Store plugin_instance reference in parameters for worker thread OAuth token access
+            parameters._plugin_instance = self._plugin_instance  # pylint: disable=protected-access
             try:
                 self._plugin_instance._configuration_parameters = parameters
 

omnata_plugin_runtime/threading_utils.py

@@ -0,0 +1,27 @@
+"""
+Utilities for thread management in the plugin runtime.
+"""
+import threading
+
+# Thread-local storage to track if we're in a managed worker thread
+# This is more reliable than checking thread names
+_thread_local = threading.local()
+
+
+def is_managed_worker_thread() -> bool:
+    """
+    Check if the current thread is a managed worker thread.
+    Returns True if running in a @managed_inbound_processing or @managed_outbound_processing worker.
+    
+    This is set by the decorator worker functions and is more reliable than checking thread names.
+    """
+    return getattr(_thread_local, 'is_managed_worker', False)
+
+
+def set_managed_worker_thread(is_worker: bool):
+    """
+    Set the flag indicating whether the current thread is a managed worker thread.
+    
+    This should only be called by the managed processing decorator worker functions.
+    """
+    _thread_local.is_managed_worker = is_worker

omnata_plugin_runtime-0.12.2a347.dist-info/METADATA

@@ -0,0 +1,56 @@
+Metadata-Version: 2.4
+Name: omnata-plugin-runtime
+Version: 0.12.2a347
+Summary: Classes and common runtime components for building and running Omnata Plugins
+License-File: LICENSE
+Author: James Weakley
+Author-email: james.weakley@omnata.com
+Requires-Python: >=3.10,<=3.13
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: annotated-types (<=0.6.0)
+Requires-Dist: certifi (<=2025.1.31)
+Requires-Dist: cffi (<=2.0.0)
+Requires-Dist: charset-normalizer (<=3.4.4)
+Requires-Dist: cryptography (<=46.0.3)
+Requires-Dist: filelock (<=3.20.0)
+Requires-Dist: idna (<=3.11)
+Requires-Dist: jinja2 (>=3.1.2,<=3.1.6)
+Requires-Dist: markupsafe (<=3.0.2)
+Requires-Dist: numpy (<=2.3.5)
+Requires-Dist: opentelemetry-api (<=1.38.0)
+Requires-Dist: packaging (<=25.0)
+Requires-Dist: pandas (<=2.3.3)
+Requires-Dist: platformdirs (<=4.5.0)
+Requires-Dist: protobuf (<=6.33.0)
+Requires-Dist: pyarrow (<=21.0.0)
+Requires-Dist: pycparser (<=2.23)
+Requires-Dist: pydantic (>=2,<=2.12.4)
+Requires-Dist: pydantic-core (<=2.41.5)
+Requires-Dist: pyjwt (<=2.10.1)
+Requires-Dist: pyopenssl (<=225.3.0)
+Requires-Dist: pytz (<=2025.2)
+Requires-Dist: pyyaml (<=6.0.3)
+Requires-Dist: requests (>=2,<=2.32.5)
+Requires-Dist: setuptools (<=80.9.0)
+Requires-Dist: snowflake-connector-python (>=3.0.0,<=4.2.0)
+Requires-Dist: snowflake-snowpark-python (>=1.20.0,<=1.44.0)
+Requires-Dist: snowflake-telemetry-python (<=0.5.0)
+Requires-Dist: tenacity (>=8,<9)
+Requires-Dist: tomlkit (<=0.13.3)
+Requires-Dist: urllib3 (<=2.5.0)
+Requires-Dist: wheel (<=0.45.1)
+Requires-Dist: wrapt (<=2.0.1)
+Description-Content-Type: text/markdown
+
+# omnata-plugin-runtime
+This package is a runtime dependency for [Omnata Plugins](https://docs.omnata.com/omnata-product-documentation/omnata-sync-for-snowflake/plugins).
+
+It contains data classes, interfaces and application logic used to perform plugin operations.
+
+For instructions on creating plugins, visit our [docs site](https://docs.omnata.com/omnata-product-documentation/omnata-sync-for-snowflake/plugins/creating-plugins).
+
+

omnata_plugin_runtime-0.12.2a347.dist-info/RECORD

@@ -0,0 +1,14 @@
+omnata_plugin_runtime/__init__.py,sha256=MS9d1whnfT_B3-ThqZ7l63QeC_8OEKTuaYV5wTwRpBA,1576
+omnata_plugin_runtime/api.py,sha256=5gbjbnFy72Xjf0E3kbG23G0V2J3CorvD5kpBn_BkdlI,8084
+omnata_plugin_runtime/configuration.py,sha256=-tN0yztdi-trgzKhLxSsPU0Ar3EBOmsNNLWFIIg3Bbc,49714
+omnata_plugin_runtime/forms.py,sha256=Lrbr3otsFDrvHWJw7v-slsW4PvEHJ6BG1Yl8oaJfiDo,20529
+omnata_plugin_runtime/json_schema.py,sha256=Wu0rByO8pFSZ3ugKqfs_yWMU24PwiC2jmoO83n9fycM,59852
+omnata_plugin_runtime/logging.py,sha256=qUtRA9syQNnjfJZHA2W18K282voXX6vHwrBIPOBo1n8,4521
+omnata_plugin_runtime/omnata_plugin.py,sha256=AJZFi9PIykFNKCWsM9mZgmXss4JyI5VWDPe2m4jnYqA,148592
+omnata_plugin_runtime/plugin_entrypoints.py,sha256=9vN1m0w7-z3qu9up1qZokfncvJlQL6tYPh5ASAuY5VQ,33023
+omnata_plugin_runtime/rate_limiting.py,sha256=qpr5esU4Ks8hMzuMpSR3gLFdor2ZUXYWCjmsQH_K6lQ,25882
+omnata_plugin_runtime/threading_utils.py,sha256=fqlKLCPTEPVYdMinf8inPKLYxwD4d4WWVMLB3a2mNqk,906
+omnata_plugin_runtime-0.12.2a347.dist-info/METADATA,sha256=8g-oocptHhzfjF_DyFkNfyVZSQus3JeA5NdofJlD99U,2235
+omnata_plugin_runtime-0.12.2a347.dist-info/WHEEL,sha256=kJCRJT_g0adfAJzTx2GUMmS80rTJIVHRCfG0DQgLq3o,88
+omnata_plugin_runtime-0.12.2a347.dist-info/licenses/LICENSE,sha256=rGaMQG3R3F5-JGDp_-rlMKpDIkg5n0SI4kctTk8eZSI,56
+omnata_plugin_runtime-0.12.2a347.dist-info/RECORD,,

{omnata_plugin_runtime-0.11.6.dist-info → omnata_plugin_runtime-0.12.2a347.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.2.1
+Generator: poetry-core 2.3.1
 Root-Is-Purelib: true
 Tag: py3-none-any

omnata_plugin_runtime-0.11.6.dist-info/METADATA

@@ -1,55 +0,0 @@
-Metadata-Version: 2.4
-Name: omnata-plugin-runtime
-Version: 0.11.6
-Summary: Classes and common runtime components for building and running Omnata Plugins
-License-File: LICENSE
-Author: James Weakley
-Author-email: james.weakley@omnata.com
-Requires-Python: >=3.9,<=3.11
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: annotated-types (<=0.6.0)
-Requires-Dist: certifi (<=2024.8.30)
-Requires-Dist: cffi (<=1.16.0)
-Requires-Dist: charset-normalizer (<=3.3.2)
-Requires-Dist: cryptography (<=43.0.0)
-Requires-Dist: filelock (<=3.13.1)
-Requires-Dist: idna (<=3.7)
-Requires-Dist: jinja2 (>=3.1.2,<=3.1.4)
-Requires-Dist: markupsafe (<=2.1.3)
-Requires-Dist: numpy (<=2.1.3)
-Requires-Dist: opentelemetry-api (<=1.23.0)
-Requires-Dist: packaging (<=24.1)
-Requires-Dist: pandas (<=2.2.3)
-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)
-Requires-Dist: pydantic-core (<=2.21.0)
-Requires-Dist: pyjwt (<=2.8.0)
-Requires-Dist: pyopenssl (<=24.2.1)
-Requires-Dist: pytz (<=2024.1)
-Requires-Dist: pyyaml (<=6.0.1)
-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.20.0,<=1.24.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
-This package is a runtime dependency for [Omnata Plugins](https://docs.omnata.com/omnata-product-documentation/omnata-sync-for-snowflake/plugins).
-
-It contains data classes, interfaces and application logic used to perform plugin operations.
-
-For instructions on creating plugins, visit our [docs site](https://docs.omnata.com/omnata-product-documentation/omnata-sync-for-snowflake/plugins/creating-plugins).
-
-

omnata_plugin_runtime-0.11.6.dist-info/RECORD

@@ -1,13 +0,0 @@
-omnata_plugin_runtime/__init__.py,sha256=MS9d1whnfT_B3-ThqZ7l63QeC_8OEKTuaYV5wTwRpBA,1576
-omnata_plugin_runtime/api.py,sha256=5gbjbnFy72Xjf0E3kbG23G0V2J3CorvD5kpBn_BkdlI,8084
-omnata_plugin_runtime/configuration.py,sha256=SffokJfgvy6V3kUsoEjXcK3GdNgHo6U3mgBEs0qBv4I,46972
-omnata_plugin_runtime/forms.py,sha256=Lrbr3otsFDrvHWJw7v-slsW4PvEHJ6BG1Yl8oaJfiDo,20529
-omnata_plugin_runtime/json_schema.py,sha256=ZfHMG-XSJBE9Smt33Y6GPpl5skF7pB1TRCf9AvWuw-Y,59705
-omnata_plugin_runtime/logging.py,sha256=qUtRA9syQNnjfJZHA2W18K282voXX6vHwrBIPOBo1n8,4521
-omnata_plugin_runtime/omnata_plugin.py,sha256=GTFxri5FIkZhOI70FQpeLoKRo3zOQlqo2dUUUTaDppE,143134
-omnata_plugin_runtime/plugin_entrypoints.py,sha256=_1pDLov3iQorGmfcae8Sw2bVjxw1vYeowBaKKNzRclQ,32629
-omnata_plugin_runtime/rate_limiting.py,sha256=qpr5esU4Ks8hMzuMpSR3gLFdor2ZUXYWCjmsQH_K6lQ,25882
-omnata_plugin_runtime-0.11.6.dist-info/METADATA,sha256=u_aYClOOFjKwP3sbJEA0Dp-YEX-p0szbOCYEPkT7tqA,2179
-omnata_plugin_runtime-0.11.6.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-omnata_plugin_runtime-0.11.6.dist-info/licenses/LICENSE,sha256=rGaMQG3R3F5-JGDp_-rlMKpDIkg5n0SI4kctTk8eZSI,56
-omnata_plugin_runtime-0.11.6.dist-info/RECORD,,