omnata-plugin-runtime 0.5.7a132__py3-none-any.whl → 0.5.7a134__py3-none-any.whl

omnata_plugin_runtime/omnata_plugin.py

@@ -38,6 +38,7 @@ import jinja2
 import pandas
 from pydantic_core import to_jsonable_python
 from pydantic import Field, TypeAdapter, ValidationError, create_model, root_validator, BaseModel
+from pydantic.dataclasses import dataclass
 from dateutil.parser import parse
 from jinja2 import Environment
 from snowflake.connector.pandas_tools import write_pandas
@@ -2228,7 +2229,25 @@ def get_nested_value(nested_dict:Dict, keys:List[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]]):
+
+@dataclass
+class SnowflakeFunctionParameter:
+    arg_name: str
+    description: str
+    arg_data_type: str
+    has_default_value: bool = False
+    default_value_clause: Optional[str] = None
+
+@dataclass
+class SnowflakeUDTFResultColumn:
+    column_name: str
+    column_data_type: str
+    description: str
+
+def consumer_udtf(
+        description: str,
+        params: List[SnowflakeFunctionParameter],
+        result_columns: List[SnowflakeUDTFResultColumn]):
     def class_decorator(cls):
         # Get the original 'process' method from the class
         cls._is_omnata_consumer_udtf = True
@@ -2242,21 +2261,14 @@ def consumer_udtf(param_docs: Dict[str, Dict[str, Any]]):
 
         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 != 'cls' and name != 'connection_parameters'
-        }
+        
+        if params[list(params.keys())[1]].annotation != Dict:
+            raise ValueError("The second argument must be a dict.")
 
         @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.")
 
             # convert the connection parameters dictionary to a ConnectionConfigurationParameters object which includes the real secrets
             if 'other_secrets_name' in connection_parameters:
@@ -2282,7 +2294,7 @@ def consumer_udtf(param_docs: Dict[str, Dict[str, Any]]):
     
     return class_decorator
 
-def find_classes_with_attribute(attribute_name: str,path:str = '.'):
+def find_consumer_udtf_classes(path:str = '.'):
     # Get the directory's absolute path
     current_dir = os.path.abspath(path)
 
@@ -2297,11 +2309,15 @@ def find_classes_with_attribute(attribute_name: str,path:str = '.'):
         # Iterate over all members of the module
         for name, obj in inspect.getmembers(module, inspect.isclass):
             # Check if the class has the specified attribute
-            if hasattr(obj, attribute_name):
+            if hasattr(obj, '_is_omnata_consumer_udtf'):
                 matching_classes.append(obj)
+    return matching_classes
 
 
-def consumer_udf(param_docs: Dict[str, Dict[str, Any]]):
+def consumer_udf(
+        description: str,
+        params: List[SnowflakeFunctionParameter],
+        result_data_type: str):
     def decorator(func):
         sig = signature(func)
         params = sig.parameters
@@ -2309,14 +2325,14 @@ def consumer_udf(param_docs: Dict[str, Dict[str, Any]]):
         # Ensure the first argument is mandatory and positional
         if list(params.keys())[0] != 'connection_parameters':
             raise ValueError("The first argument should be 'connection_parameters'.")
+        if params[list(params.keys())[0]].annotation != Dict:
+            raise ValueError("The first argument must be a dict.")
 
         @wraps(func)
         def wrapper(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.")
 
             # convert the connection parameters dictionary to a ConnectionConfigurationParameters object which includes the real secrets
             if 'other_secrets_name' in connection_parameters:
@@ -2340,3 +2356,23 @@ def consumer_udf(param_docs: Dict[str, Dict[str, Any]]):
         return wrapper
     
     return decorator
+
+def find_consumer_udf_functions(path:str = '.'):
+    # Get the current directory's absolute path
+    current_dir = os.path.abspath(path)
+
+    # List to hold the classes that match the attribute
+    matching_classes = []
+
+    # Iterate over all modules in the current directory
+    for _, module_name, _ in pkgutil.iter_modules([current_dir]):
+        # Import the module
+        module = importlib.import_module(module_name)
+
+        # Iterate over all members of the module
+        for name, obj in inspect.getmembers(module, inspect.isfunction):
+            # Check if the class has the specified attribute
+            if hasattr(obj, '_is_omnata_consumer_udf'):
+                matching_classes.append(obj)
+
+    return matching_classes

{omnata_plugin_runtime-0.5.7a132.dist-info → omnata_plugin_runtime-0.5.7a134.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: omnata-plugin-runtime
-Version: 0.5.7a132
+Version: 0.5.7a134
 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.7a132.dist-info → omnata_plugin_runtime-0.5.7a134.dist-info}/RECORD

@@ -3,10 +3,10 @@ omnata_plugin_runtime/api.py,sha256=FxzTqri4no8ClkOm7vZADG8aD47jcGBCTTQDEORmOJM,
 omnata_plugin_runtime/configuration.py,sha256=TI6GaVFhewVawBCaYN34GujY57qEP6q2nik4YpSEk5s,38100
 omnata_plugin_runtime/forms.py,sha256=GzSPEwcijsoPCXEO1mHiE8ylvX_KSE5TkhwqkymA2Ss,19755
 omnata_plugin_runtime/logging.py,sha256=bn7eKoNWvtuyTk7RTwBS9UARMtqkiICtgMtzq3KA2V0,3272
-omnata_plugin_runtime/omnata_plugin.py,sha256=eV1QKOqdxktkVB3ZEQPS7A9a9q7UOsyEhfyZ-JwV5iw,117231
+omnata_plugin_runtime/omnata_plugin.py,sha256=WgdUO5PvwwhekdB-hSLr_4hqjK--BXmm6G3opgeDMA8,118075
 omnata_plugin_runtime/plugin_entrypoints.py,sha256=PFSLsYEVnWHVvSoOYTtTK2JY6pp6_8_eYP53WqLRiPE,27975
 omnata_plugin_runtime/rate_limiting.py,sha256=DVQ_bc-mVLBkrU1PTns1MWXhHiLpSB5HkWCcdePtJ2A,25611
-omnata_plugin_runtime-0.5.7a132.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-omnata_plugin_runtime-0.5.7a132.dist-info/METADATA,sha256=YHhi631jWmg1ZBJuNsry--DCdhMxWVNynL2ZreqnS7w,1985
-omnata_plugin_runtime-0.5.7a132.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-omnata_plugin_runtime-0.5.7a132.dist-info/RECORD,,
+omnata_plugin_runtime-0.5.7a134.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
+omnata_plugin_runtime-0.5.7a134.dist-info/METADATA,sha256=JGVpgMpr1s7vVOCc9qdKTNKKLxoFnOjx27LtePTLDmM,1985
+omnata_plugin_runtime-0.5.7a134.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+omnata_plugin_runtime-0.5.7a134.dist-info/RECORD,,