omnata-plugin-runtime 0.5.6a128__tar.gz → 0.5.7a130__tar.gz

{omnata_plugin_runtime-0.5.6a128 → omnata_plugin_runtime-0.5.7a130}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omnata-plugin-runtime
-Version: 0.5.6a128
+Version: 0.5.7a130
 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.5.6a128 → omnata_plugin_runtime-0.5.7a130}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "omnata-plugin-runtime"
-version = "0.5.6-a128"
+version = "0.5.7-a130"
 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.5.6a128 → omnata_plugin_runtime-0.5.7a130}/src/omnata_plugin_runtime/configuration.py

@@ -3,12 +3,14 @@ Omnata Plugin Runtime configuration objects.
 Includes data container classes related to plugin configuration.
 """
 from __future__ import annotations
+import json
 import sys
+import logging
 from typing import Any, List, Dict, Literal, Union, Optional
 from enum import Enum
 
 from abc import ABC
-from pydantic import BaseModel, Field, PrivateAttr, SerializationInfo, field_validator, model_serializer, validator  # pylint: disable=no-name-in-module
+from pydantic import BaseModel, Field, PrivateAttr, SerializationInfo, TypeAdapter, field_validator, model_serializer, validator  # pylint: disable=no-name-in-module
 
 if tuple(sys.version_info[:2]) >= (3, 9):
     # Python 3.9 and above
@@ -17,6 +19,8 @@ else:
     # Python 3.8 and below
     from typing_extensions import Annotated
 
+logger = logging.getLogger(__name__)
+
 
 class MapperType(str, Enum):
     FIELD_MAPPING_SELECTOR = "field_mapping_selector"
@@ -860,3 +864,29 @@ StoredStreamConfiguration.model_rebuild()
 InboundSyncStreamsConfiguration.model_rebuild()
 StoredFieldMappings.model_rebuild()
 OutboundSyncConfigurationParameters.model_rebuild()
+
+
+def get_secrets(oauth_secret_name: Optional[str], other_secrets_name: Optional[str]
+) -> Dict[str, StoredConfigurationValue]:
+    connection_secrets = {}
+    import _snowflake # pylint: disable=import-error, import-outside-toplevel # type: ignore
+    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
+    return connection_secrets
+    

{omnata_plugin_runtime-0.5.6a128 → omnata_plugin_runtime-0.5.7a130}/src/omnata_plugin_runtime/omnata_plugin.py

@@ -4,6 +4,7 @@ Omnata Plugin Runtime.
 Includes data container classes and defines the contract for a plugin.
 """
 from __future__ import annotations
+from inspect import signature
 import sys
 from typing import Union
 if tuple(sys.version_info[:2]) >= (3, 9):
@@ -31,7 +32,7 @@ from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Type,
 import jinja2
 import pandas
 from pydantic_core import to_jsonable_python
-from pydantic import Field, TypeAdapter, root_validator, BaseModel
+from pydantic import Field, TypeAdapter, ValidationError, create_model, root_validator, BaseModel
 from dateutil.parser import parse
 from jinja2 import Environment
 from snowflake.connector.pandas_tools import write_pandas
@@ -65,6 +66,7 @@ from .configuration import (
     StreamConfiguration,
     SubscriptableBaseModel,
     SyncConfigurationParameters,
+    get_secrets
 )
 from .forms import (
     ConnectionMethod,
@@ -2220,3 +2222,74 @@ def get_nested_value(nested_dict:Dict, keys:List[str]):
     if isinstance(keys, str):
         keys = [keys]
     return reduce(lambda d, key: d.get(key) if isinstance(d, dict) else None, keys, nested_dict)
+
+def consumer_udtf(param_docs: Dict[str, Dict[str, Any]]):
+    def class_decorator(cls):
+        # Get the original 'process' method from the class
+        cls._is_omnata_consumer_udtf = True
+        original_process = getattr(cls, 'process')
+        sig = signature(original_process)
+        params = sig.parameters
+
+        # Ensure the first argument is mandatory and positional
+        if list(params.keys())[0] != 'self':
+            raise ValueError("The first argument should be 'self' for class methods.")
+
+        if list(params.keys())[1] != 'connection_parameters':
+            raise ValueError("The second argument should be 'connection_parameters'.")
+
+        # Create a Pydantic model based on the function signature and the provided documentation
+        fields = {
+            name: (param.annotation, param_docs.get(name, {}).get("default", param.default))
+            for name, param in params.items()
+            if name != 'self' and name != 'mandatory_arg'
+        }
+
+        DynamicModel = create_model('DynamicModel', **fields)
+
+        # Attach the documentation to the function
+        original_process.__doc__ = original_process.__doc__ or ""
+        original_process.__doc__ += "\n\nArgs:\n"
+        for name, param in params.items():
+            if name != 'self':
+                doc = param_docs.get(name, {})
+                original_process.__doc__ += f"    {name} ({param.annotation.__name__}): {doc.get('description', 'No description provided.')}\n"
+
+        @wraps(original_process)
+        def wrapped_process(self, connection_parameters, *args, **kwargs):
+            if connection_parameters is None:
+                raise ValueError("Connection not found")
+            
+            if not isinstance(connection_parameters, Dict):
+                raise ValueError("The first argument must be an object, the result of calling PLUGIN_CONNECTION.")
+
+            try:
+                # Validate the arguments using the dynamic Pydantic model
+                validated_args = DynamicModel(**kwargs)
+            except ValidationError as e:
+                raise ValueError(f"Argument validation error: {e}")
+            
+            # convert the connection parameters dictionary to a ConnectionConfigurationParameters object which includes the real secrets
+            if 'other_secrets_name' in connection_parameters:
+                # this is the new way, where the sync engine only passes the name of the secret
+                oauth_secrets_name = None
+                if 'oauth_secret_name' in connection_parameters:
+                    oauth_secrets_name = connection_parameters['oauth_secret_name']
+                    del connection_parameters['oauth_secret_name']
+                result = get_secrets(oauth_secrets_name,connection_parameters['other_secrets_name'])
+                connection_parameters['connection_secrets'] = result
+                del connection_parameters['other_secrets_name']
+                parameters = ConnectionConfigurationParameters.model_validate(connection_parameters)
+            else:
+                # deprecated way, where the sync engine passes the secrets directly
+                parameters = ConnectionConfigurationParameters.model_validate(connection_parameters)
+
+            # Pass the validated arguments to the function
+            return original_process(self, parameters, *args, **validated_args.dict())
+        # Replace the original 'process' method with the wrapped version
+        setattr(cls, 'process', wrapped_process)
+        
+        return cls
+    
+    return class_decorator
+

{omnata_plugin_runtime-0.5.6a128 → omnata_plugin_runtime-0.5.7a130}/src/omnata_plugin_runtime/plugin_entrypoints.py

@@ -20,6 +20,7 @@ from .configuration import (
     OutboundSyncStrategy,
     StoredConfigurationValue,
     StoredMappingValue,
+    get_secrets
 )
 from .forms import ConnectionMethod, FormInputField, FormOption
 from .logging import OmnataPluginLogHandler
@@ -71,7 +72,7 @@ class PluginEntrypoint:
     def sync(self, sync_request: Dict):
         logger.info("Entered sync method")
         request = TypeAdapter(SyncRequestPayload).validate_python(sync_request)
-        connection_secrets = self.get_secrets(
+        connection_secrets = get_secrets(
             request.oauth_secret_name, request.other_secrets_name
         )
         omnata_log_handler = OmnataPluginLogHandler(
@@ -270,7 +271,7 @@ class PluginEntrypoint:
         sync_strategy = normalise_nulls(sync_strategy)
         oauth_secret_name = normalise_nulls(oauth_secret_name)
         other_secrets_name = normalise_nulls(other_secrets_name)
-        connection_secrets = self.get_secrets(oauth_secret_name, other_secrets_name)
+        connection_secrets = get_secrets(oauth_secret_name, other_secrets_name)
         connection_parameters = TypeAdapter(
             Dict[str, StoredConfigurationValue]).validate_python(connection_parameters)
         sync_parameters = TypeAdapter(
@@ -324,7 +325,7 @@ class PluginEntrypoint:
         logger.info("Entered list_streams method")
         oauth_secret_name = normalise_nulls(oauth_secret_name)
         other_secrets_name = normalise_nulls(other_secrets_name)
-        connection_secrets = self.get_secrets(oauth_secret_name, other_secrets_name)
+        connection_secrets = get_secrets(oauth_secret_name, other_secrets_name)
         connection_parameters = TypeAdapter(
             Dict[str, StoredConfigurationValue]).validate_python(connection_parameters)
         sync_parameters = TypeAdapter(
@@ -411,31 +412,6 @@ class PluginEntrypoint:
                     raise e
         return [e.model_dump() for e in events]
 
-    def get_secrets(
-        self, oauth_secret_name: Optional[str], other_secrets_name: Optional[str]
-    ) -> Dict[str, StoredConfigurationValue]:
-        connection_secrets = {}
-        import _snowflake # pylint: disable=import-error, import-outside-toplevel # type: ignore
-        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
-        return connection_secrets
-    
     def ngrok_post_tunnel_fields(
         self,
         connection_method: str,
@@ -447,7 +423,7 @@ class PluginEntrypoint:
         logger.info("Entered ngrok_post_tunnel_fields method")
         oauth_secret_name = normalise_nulls(oauth_secret_name)
         other_secrets_name = normalise_nulls(other_secrets_name)
-        connection_secrets = self.get_secrets(oauth_secret_name, other_secrets_name)
+        connection_secrets = get_secrets(oauth_secret_name, other_secrets_name)
         connection_parameters = TypeAdapter(
             Dict[str, StoredConfigurationValue]).validate_python(connection_parameters)
         parameters = ConnectionConfigurationParameters(
@@ -493,7 +469,7 @@ class PluginEntrypoint:
     ):
         logger.info("Entered connect method")
         logger.info(f"Connection parameters: {connection_parameters}")
-        connection_secrets = self.get_secrets(oauth_secret_name, other_secrets_name)
+        connection_secrets = get_secrets(oauth_secret_name, other_secrets_name)
 
         from omnata_plugin_runtime.omnata_plugin import (
             ConnectionConfigurationParameters,
@@ -540,7 +516,7 @@ class PluginEntrypoint:
                     oauth_secret_name: Optional[str],
                     other_secrets_name: Optional[str]):
         logger.info("Entered api_limits method")
-        connection_secrets = self.get_secrets(oauth_secret_name, other_secrets_name)
+        connection_secrets = get_secrets(oauth_secret_name, other_secrets_name)
         from omnata_plugin_runtime.omnata_plugin import (
             ConnectionConfigurationParameters,
         )