omnata-plugin-runtime 0.5.7a147__tar.gz → 0.5.7a149__tar.gz

{omnata_plugin_runtime-0.5.7a147 → omnata_plugin_runtime-0.5.7a149}/PKG-INFO

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

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

@@ -2381,56 +2381,68 @@ def omnata_udtf(
     """
     def class_decorator(cls):
         # Get the original 'process' method from the class
+        if not hasattr(cls, 'process'):
+            raise ValueError("The class must have a 'process' method.")
         original_process = getattr(cls, 'process')
         sig = signature(original_process)
         function_params = sig.parameters
-        if len(function_params) < 2:
-            raise ValueError("The function must have the self parameter, plus at least the mandatory 'connection_parameters' parameter.")
-        # Ensure the first argument is mandatory and positional
+        if len(function_params) < 1:
+            raise ValueError("The 'process' function must have at least one parameter.")
+        
         first_param_name = list(function_params.keys())[0]
-
         if first_param_name != 'self':
-            raise ValueError(f"The first argument should be 'self', instead it was '{first_param_name}'.")
+            raise ValueError(f"The first argument for the 'process' function should be 'self', instead it was '{first_param_name}'.")
 
-        second_param_name = list(function_params.keys())[1]
+        cls._is_omnata_udtf = True
+        cls._omnata_udtf_name = name
+        cls._omnata_udtf_description = description
+        cls._omnata_udtf_params = params
+        cls._omnata_udtf_result_columns = result_columns
+        cls._omnata_udtf_expose_to_consumer = expose_to_consumer
 
+        if not expose_to_consumer:
+            # If not exposing to the consumer, there are no further requirements
+            return cls
+        
+        if len(function_params) < 2:
+            raise ValueError("When exposing the udtf to consumers, the 'process' function must have the self parameter, plus at least the mandatory 'connection_parameters' parameter.")
+        second_param_name = list(function_params.keys())[1]
         if second_param_name != 'connection_parameters':
             raise ValueError(f"The second argument should be 'connection_parameters', instead it was {second_param_name}.")
-
-
         if function_params[second_param_name].annotation != ConnectionConfigurationParameters:
             raise ValueError(f"The second argument must be a ConnectionConfigurationParameters, instead it was a {function_params[second_param_name].annotation}.")
 
-        @wraps(original_process)
-        def wrapped_process(self, connection_parameters, *args, **kwargs):
-            if connection_parameters is None:
-                raise ValueError("Connection not found")
+        if params[0].name.upper() != 'CONNECTION_PARAMETERS':
+            params = [SnowflakeFunctionParameter(
+                                    name='CONNECTION_PARAMETERS',
+                                    data_type='OBJECT',
+                                    description='The connection object, obtained from calling PLUGIN.PLUGIN_CONNECTION.')] + params
+            cls._omnata_udtf_params = params
+        if len(params) != len(function_params) -1:
+            raise ValueError(f"You must document all the parameters of the 'process' function in the @omnata_udtf decorator in the same order ('connection_parameters' will be included automatically).")
 
+        @wraps(original_process)
+        def wrapped_process(self, connection_parameter_arg, *args, **kwargs):
+            if connection_parameter_arg is None:
+                raise ValueError("Connection not found")           
+            
             # convert the connection parameters dictionary to a ConnectionConfigurationParameters object which includes the real secrets
-            if 'other_secrets_name' in connection_parameters:
+            if 'other_secrets_name' in connection_parameter_arg:
                 # 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)
+                if 'oauth_secret_name' in connection_parameter_arg:
+                    oauth_secrets_name = connection_parameter_arg['oauth_secret_name']
+                    del connection_parameter_arg['oauth_secret_name']
+                result = get_secrets(oauth_secrets_name,connection_parameter_arg['other_secrets_name'])
+                connection_parameter_arg['connection_secrets'] = result
+                del connection_parameter_arg['other_secrets_name']
+            
+            parameters = ConnectionConfigurationParameters.model_validate(connection_parameter_arg)
 
             # Pass the validated arguments to the function
             return original_process(self, parameters, *args, **kwargs)
         # Replace the original 'process' method with the wrapped version
         setattr(cls, 'process', wrapped_process)
-        cls._is_omnata_udtf = True
-        cls._omnata_udtf_name = name
-        cls._omnata_udtf_description = description
-        cls._omnata_udtf_params = params
-        cls._omnata_udtf_result_columns = result_columns
-        cls._omnata_udtf_expose_to_consumer = expose_to_consumer
         return cls
     
     return class_decorator
@@ -2475,11 +2487,22 @@ def omnata_udf(
         result_data_type: str,
         expose_to_consumer: bool):
     """
-    A decorator for a function which should be automatically exposed to the consumer as a UDF
+    A decorator for a function which will be created in the native application.
     """
     def decorator(func):
         sig = signature(func)
         function_params = sig.parameters
+
+        if not expose_to_consumer:
+            # If not exposing to the consumer, there are no further requirements
+            func._is_omnata_udf = True
+            func._omnata_udf_name = name
+            func._omnata_udf_description = description
+            func._omnata_udf_params = params
+            func._omnata_udf_result_data_type = result_data_type
+            func._omnata_udf_expose_to_consumer = expose_to_consumer
+            return func
+        
         if len(function_params) == 0:
             raise ValueError("The function must have at least one parameter.")
         # Ensure the first argument is mandatory and positional
@@ -2488,27 +2511,29 @@ def omnata_udf(
             raise ValueError(f"The first argument should be 'connection_parameters', instead it was '{first_param_name}'.")
         if function_params[first_param_name].annotation != ConnectionConfigurationParameters:
             raise ValueError(f"The first argument must be a ConnectionConfigurationParameters, instead it was a {function_params[first_param_name].annotation}.")
+        if params[0].name.upper() != 'CONNECTION_PARAMETERS':
+            params = [SnowflakeFunctionParameter(
+                                    name='CONNECTION_PARAMETERS',
+                                    data_type='OBJECT',
+                                    description='The connection object, obtained from calling PLUGIN.PLUGIN_CONNECTION.')] + params
+            func._omnata_udf_params = params
+        if len(params) != len(function_params):
+            raise ValueError(f"You must document all the parameters of the function in the @omnata_udf decorator in the same order ('connection_parameters' will be included automatically).")
 
         @wraps(func)
-        def wrapper(self, connection_parameters, *args, **kwargs):
-            if connection_parameters is None:
-                raise ValueError("Connection not found")
-            
-
+        def wrapper(connection_parameter_arg, *args, **kwargs):
             # convert the connection parameters dictionary to a ConnectionConfigurationParameters object which includes the real secrets
-            if 'other_secrets_name' in connection_parameters:
+            if 'other_secrets_name' in connection_parameter_arg:
                 # 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)
+                if 'oauth_secret_name' in connection_parameter_arg:
+                    oauth_secrets_name = connection_parameter_arg['oauth_secret_name']
+                    del connection_parameter_arg['oauth_secret_name']
+                result = get_secrets(oauth_secrets_name,connection_parameter_arg['other_secrets_name'])
+                connection_parameter_arg['connection_secrets'] = result
+                del connection_parameter_arg['other_secrets_name']
+                
+            parameters = ConnectionConfigurationParameters.model_validate(connection_parameter_arg)
 
             # Pass the validated arguments to the function
             return func(parameters, *args, **kwargs)