omnata-plugin-runtime 0.12.1__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

@@ -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,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):
         """
@@ -1344,18 +1426,23 @@ class InboundSyncRequest(SyncRequest):
                 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"
@@ -1374,6 +1461,10 @@ class InboundSyncRequest(SyncRequest):
             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:
             self._apply_results_if_size_exceeded()
@@ -1450,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(
@@ -2184,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()}"
@@ -2323,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")
 
@@ -2367,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
@@ -2505,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")
 
@@ -2762,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
     """
@@ -2778,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.1.dist-info → omnata_plugin_runtime-0.12.2a347.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: omnata-plugin-runtime
-Version: 0.12.1
+Version: 0.12.2a347
 Summary: Classes and common runtime components for building and running Omnata Plugins
 License-File: LICENSE
 Author: James Weakley
@@ -36,8 +36,8 @@ 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,<4)
-Requires-Dist: snowflake-snowpark-python (>=1.20.0,<=1.43.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)

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.12.1.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.12.1.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.1.dist-info/METADATA,sha256=SCl6ee1e3Q8DN0wa47snuMAOBABw387hC54HXuYSTcs,2222
-omnata_plugin_runtime-0.12.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-omnata_plugin_runtime-0.12.1.dist-info/licenses/LICENSE,sha256=rGaMQG3R3F5-JGDp_-rlMKpDIkg5n0SI4kctTk8eZSI,56
-omnata_plugin_runtime-0.12.1.dist-info/RECORD,,