omnata-plugin-runtime 0.12.1a330__py3-none-any.whl → 0.12.2a337__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
     )
+    _sync_request: Optional[Any] = PrivateAttr(  # Reference to SyncRequest for worker thread access
+        default=None
+    )
 
     @model_validator(mode='after')
     def validate_ngrok_tunnel_settings(self) -> Self:
@@ -739,6 +742,23 @@ 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._sync_request is not None:
+                logger.debug(f"Worker thread requesting access_token via secrets service")
+                try:
+                    secrets = self._sync_request.request_secrets_from_main_thread(
+                        self.access_token_secret_name, None
+                    )
+                    if 'access_token' in secrets:
+                        return secrets['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,10 +1025,34 @@ 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), it will automatically
+    route the request through the secrets service to avoid threading issues with _snowflake.get_oauth_access_token.
+    
+    :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
+    
+    # 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 sync_request is not None:
+        logger.debug(f"Worker thread requesting secrets via secrets service")
+        try:
+            return sync_request.request_secrets_from_main_thread(oauth_secret_name, other_secrets_name)
+        except Exception as e:
+            logger.error(f"Error requesting secrets from main thread: {e}")
+            raise
+    
+    # Otherwise, call _snowflake functions directly (main thread)
     if oauth_secret_name is not None:
         connection_secrets["access_token"] = StoredConfigurationValue(
             value=_snowflake.get_oauth_access_token(oauth_secret_name)

omnata_plugin_runtime/omnata_plugin.py

@@ -101,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"]
 
@@ -375,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:
@@ -499,6 +505,105 @@ class SyncRequest(ABC):
             cancellation_token.wait(20)
         logger.info("cancel checking worker exiting")
 
+    def _service_secrets_requests(self):
+        """
+        Services any pending secrets 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
+        from .configuration import StoredConfigurationValue
+        
+        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')
+            other_secrets_name = request.get('other_secrets_name')
+            response_queue = request['response_queue']
+            
+            logger.debug(f"Main thread servicing secrets request")
+            
+            try:
+                # Call _snowflake functions directly (we're on the main thread now)
+                connection_secrets = {}
+                if oauth_secret_name is not None:
+                    connection_secrets["access_token"] = StoredConfigurationValue(
+                        value=_snowflake.get_oauth_access_token(oauth_secret_name)
+                    )
+                if other_secrets_name is not None:
+                    try:
+                        secret_string_content = _snowflake.get_generic_secret_string(
+                            other_secrets_name
+                        )
+                        if len(secret_string_content) > 2:
+                            other_secrets = json.loads(secret_string_content)
+                            connection_secrets = {
+                                **connection_secrets,
+                                **TypeAdapter(Dict[str, StoredConfigurationValue]).validate_python(other_secrets),
+                            }
+                    except Exception as exception:
+                        logger.error(f"Error parsing secrets content for secret {other_secrets_name}: {str(exception)}")
+                        raise ValueError(f"Error parsing secrets content: {str(exception)}") from exception
+                
+                # Put the result in the response queue for the requesting thread
+                response_queue.put({
+                    'success': True,
+                    'result': connection_secrets
+                })
+            except Exception as e:
+                logger.error(f"Error servicing secrets request: {e}")
+                response_queue.put({
+                    'success': False,
+                    'error': str(e)
+                })
+            finally:
+                self._secrets_request_queue.task_done()
+        
+        return serviced_any
+
+    def request_secrets_from_main_thread(self, oauth_secret_name: Optional[str], 
+                                        other_secrets_name: Optional[str],
+                                        timeout: int = 30) -> Dict[str, Any]:
+        """
+        Request secrets from the main thread. This should be called from worker threads
+        when they need to access secrets 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 other_secrets_name: The name of other secrets to retrieve
+        :param timeout: Maximum time to wait for the response in seconds
+        :return: Dictionary of StoredConfigurationValue objects
+        :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 secrets from main thread")
+        
+        # Put the request in the queue with its own response queue
+        self._secrets_request_queue.put({
+            'oauth_secret_name': oauth_secret_name,
+            'other_secrets_name': other_secrets_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 secrets: {response['error']}")
+        except queue.Empty:
+            raise TimeoutError(f"Timeout waiting for secrets request after {timeout} seconds")
+
     @abstractmethod
     def apply_results_queue(self):
         """
@@ -2184,6 +2289,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()}"
@@ -2323,6 +2431,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 secrets requests from worker threads while we wait
+                self._sync_request._service_secrets_requests()
                 time.sleep(1)  # Avoid busy waiting
             logger.info("All workers completed processing")
 
@@ -2367,6 +2477,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
@@ -2505,6 +2618,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 secrets requests from worker threads while we wait
+                    self._sync_request._service_secrets_requests()
                     time.sleep(1)  # Avoid busy waiting
                 logger.info("All workers completed processing")
 
@@ -2762,7 +2877,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
     """
@@ -2778,6 +2893,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 sync_request reference in parameters for worker thread access
+            parameters._sync_request = outbound_sync_request  # 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 sync_request reference in parameters for worker thread access
+            parameters._sync_request = inbound_sync_request  # 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.1a330.dist-info → omnata_plugin_runtime-0.12.2a337.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: omnata-plugin-runtime
-Version: 0.12.1a330
+Version: 0.12.2a337
 Summary: Classes and common runtime components for building and running Omnata Plugins
 License-File: LICENSE
 Author: James Weakley

omnata_plugin_runtime-0.12.2a337.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=-y9TmhDz6iI6u1AZ0sPG7HIJLOpiBpuscUXFQm7SkRs,49351
+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=FLVU88CjTwI52OOxAsouu1DZh-stY_b-r1uc7tgMTn8,149390
+omnata_plugin_runtime/plugin_entrypoints.py,sha256=0glIkuQpva-C-a3Cn3KNI6Jj2u7fpdWV1aiUgrEz3Mo,32986
+omnata_plugin_runtime/rate_limiting.py,sha256=qpr5esU4Ks8hMzuMpSR3gLFdor2ZUXYWCjmsQH_K6lQ,25882
+omnata_plugin_runtime/threading_utils.py,sha256=fqlKLCPTEPVYdMinf8inPKLYxwD4d4WWVMLB3a2mNqk,906
+omnata_plugin_runtime-0.12.2a337.dist-info/METADATA,sha256=eQIYDh_p1MuA0hE6rQeZaTdWV9VL5TLWk1A1j6n5jvs,2226
+omnata_plugin_runtime-0.12.2a337.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+omnata_plugin_runtime-0.12.2a337.dist-info/licenses/LICENSE,sha256=rGaMQG3R3F5-JGDp_-rlMKpDIkg5n0SI4kctTk8eZSI,56
+omnata_plugin_runtime-0.12.2a337.dist-info/RECORD,,

omnata_plugin_runtime-0.12.1a330.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=8FT3XNdZzty76OldvcxdKpbKrPENKjAIbwa_rxceVyg,143564
-omnata_plugin_runtime/plugin_entrypoints.py,sha256=_1pDLov3iQorGmfcae8Sw2bVjxw1vYeowBaKKNzRclQ,32629
-omnata_plugin_runtime/rate_limiting.py,sha256=qpr5esU4Ks8hMzuMpSR3gLFdor2ZUXYWCjmsQH_K6lQ,25882
-omnata_plugin_runtime-0.12.1a330.dist-info/METADATA,sha256=S5v-G348UEulYX6Uai-HY9xuEBqGq6Ija1SPwgiKoTU,2226
-omnata_plugin_runtime-0.12.1a330.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-omnata_plugin_runtime-0.12.1a330.dist-info/licenses/LICENSE,sha256=rGaMQG3R3F5-JGDp_-rlMKpDIkg5n0SI4kctTk8eZSI,56
-omnata_plugin_runtime-0.12.1a330.dist-info/RECORD,,