omnata-plugin-runtime 0.5.7a152__tar.gz → 0.5.7a153__tar.gz

{omnata_plugin_runtime-0.5.7a152 → omnata_plugin_runtime-0.5.7a153}/PKG-INFO

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

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

@@ -39,7 +39,6 @@ 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
@@ -2231,8 +2230,7 @@ def get_nested_value(nested_dict:Dict, keys:List[str]):
     return reduce(lambda d, key: d.get(key) if isinstance(d, dict) else None, keys, nested_dict)
 
 
-@dataclass
-class SnowflakeFunctionParameter:
+class SnowflakeFunctionParameter(BaseModel):
     """
     Represents a parameter for a Snowflake UDF or UDTF
     """
@@ -2246,8 +2244,7 @@ class SnowflakeFunctionParameter:
             return f"{self.name} {self.data_type} default {self.default_value_clause}"
         return f"{self.name} {self.data_type}"
 
-@dataclass
-class SnowflakeUDTFResultColumn:
+class SnowflakeUDTFResultColumn(BaseModel):
     """
     Represents a result column for a Snowflake UDTF
     """
@@ -2256,49 +2253,21 @@ class SnowflakeUDTFResultColumn:
     def __str__(self):
         return f"{self.name} {self.data_type}"
 
-@dataclass
-class PythonUDTFDefinition:
+class UDTFDefinition(BaseModel):
     """
     The information needed by the plugin uploader to put a Python UDTF definition into the setup script.
     Do not use this class directly in plugins, instead use the omnata_udtf decorator.
     """
-    name: str
-    description: str
-    params: List[SnowflakeFunctionParameter]
-    result_columns: List[SnowflakeUDTFResultColumn]
-    handler: str
-    expose_to_consumer: bool
-    packages: Optional[List[str]] = None
-
-    def __str__(self):
-        param_str = ', '.join([str(param) for param in self.params])
-        table_result_columns = ', '.join([f"{col.name} {col.data_type}" for col in self.result_columns])
-        packages_str = ', '.join([f"'{p}'" for p in self.packages])
-        return f"""CREATE OR REPLACE FUNCTION UDFS.{self.name}({param_str})
-RETURNS TABLE({table_result_columns})
-LANGUAGE PYTHON
-RUNTIME_VERSION=3.10
-COMMENT = $${self.description}$$
-PACKAGES = ({packages_str})
-IMPORTS = ('/app.zip')
-HANDLER='{self.handler}';
-"""
-    
-@dataclass
-class JavaUDTFDefinition:
-    """
-    The information needed by the plugin uploader to put a Java UDTF definition into the setup script.
-    These are created directly in python, since we don't want the omnata cli to need to do any Java reflection.
-    Any .jar files placed in the udf_direct_imports directory will be uploaded to the stage and can be
-    added to the imports parameter as (e.g.) '/my_jar.jar'.
-    """
-    name: str
-    description: str
-    params: List[SnowflakeFunctionParameter]
-    result_columns: List[SnowflakeUDTFResultColumn]
-    packages: List[str]
-    imports: List[str]
-    handler: str
+    name: str = Field(..., title="The name of the UDTF")
+    language: Literal['python','java'] = Field(..., title="The language of the UDF")
+    runtime_version: str = Field(..., title="The runtime version of the UDF (language dependent)")
+    description: str = Field(..., title="A description of the UDTF")
+    params: List[SnowflakeFunctionParameter] = Field(..., title="The parameters of the UDTF")
+    result_columns: List[SnowflakeUDTFResultColumn] = Field(..., title="The result columns of the UDTF")
+    handler: str = Field(..., title="The handler class/function for the UDTF")
+    expose_to_consumer: bool = Field(..., title="Whether the UDTF should be exposed to consumers")
+    imports: Optional[List[str]] = Field(None, title="A list of imports required by the UDF")
+    packages: Optional[List[str]] = Field(None, title="A list of packages required by the UDTF")
 
     def __str__(self):
         param_str = ', '.join([str(param) for param in self.params])
@@ -2307,54 +2276,29 @@ class JavaUDTFDefinition:
         imports_str = ', '.join([f"'{i}'" for i in self.imports])
         return f"""CREATE OR REPLACE FUNCTION UDFS.{self.name}({param_str})
 RETURNS TABLE({table_result_columns})
-LANGUAGE PYTHON
-RUNTIME_VERSION=3.10
+LANGUAGE {self.language.upper()}
+RUNTIME_VERSION={self.runtime_version}
 COMMENT = $${self.description}$$
 PACKAGES = ({packages_str})
 IMPORTS = ({imports_str})
 HANDLER='{self.handler}';
 """
 
-@dataclass
-class PythonUDFDefinition:
+class UDFDefinition(BaseModel):
     """
     The information needed by the plugin uploader to put a Python UDF definition into the setup script.
     Do not use this class directly in plugins, instead use the omnata_udf decorator.
     """
-    name: str
-    description: str
-    params: List[SnowflakeFunctionParameter]
-    result_data_type: str
-    handler: str
-    expose_to_consumer: bool
-    packages: Optional[List[str]] = None
-
-    def __str__(self):
-        param_str = ', '.join([str(param) for param in self.params])
-        packages_str = ', '.join([f"'{p}'" for p in self.packages])
-        return f"""CREATE OR REPLACE FUNCTION UDFS.{self.name}({param_str})
-RETURNS {self.result_data_type}
-LANGUAGE PYTHON
-RUNTIME_VERSION=3.10
-COMMENT = $${self.description}$$
-PACKAGES = ({packages_str})
-IMPORTS = ('/app.zip')
-HANDLER='{self.handler}';
-"""
-
-@dataclass
-class JavaUDFDefinition:
-    """
-    The information needed by the plugin uploader to put a java UDF definition into the setup script.
-    """
-    name: str
-    description: str
-    params: List[SnowflakeFunctionParameter]
-    result_data_type: str
-    packages: List[str]
-    imports: List[str]
-    handler: str
-    expose_to_consumer: bool
+    name: str = Field(..., title="The name of the UDF")
+    language: Literal['python','java'] = Field(..., title="The language of the UDF")
+    runtime_version: str = Field(..., title="The runtime version of the UDF (language dependent)")
+    description: str = Field(..., title="A description of the UDF")
+    params: List[SnowflakeFunctionParameter] = Field(..., title="The parameters of the UDF")
+    result_data_type: str = Field(..., title="The data type returned by the UDF")
+    handler: str = Field(..., title="The handler class/function for the UDF")
+    expose_to_consumer: bool = Field(..., title="Whether the UDF should be exposed to consumers")
+    imports: Optional[List[str]] = Field(None, title="A list of imports required by the UDF")
+    packages: Optional[List[str]] = Field(None, title="A list of packages required by the UDF")
 
     def __str__(self):
         param_str = ', '.join([str(param) for param in self.params])
@@ -2362,8 +2306,8 @@ class JavaUDFDefinition:
         imports_str = ', '.join([f"'{i}'" for i in self.imports])
         return f"""CREATE OR REPLACE FUNCTION UDFS.{self.name}({param_str})
 RETURNS {self.result_data_type}
-LANGUAGE PYTHON
-RUNTIME_VERSION=3.10
+LANGUAGE {self.language.upper()}
+RUNTIME_VERSION={self.runtime_version}
 COMMENT = $${self.description}$$
 PACKAGES = ({packages_str})
 IMPORTS = ({imports_str})
@@ -2447,7 +2391,7 @@ def omnata_udtf(
     
     return class_decorator
 
-def find_udtf_classes(path:str = '.') -> List[PythonUDTFDefinition | JavaUDTFDefinition]:
+def find_udtf_classes(path:str = '.') -> List[UDTFDefinition]:
     """
     Finds all classes in the specified directory which have the 'omnata_udtf' decorator applied
     """
@@ -2466,15 +2410,16 @@ def find_udtf_classes(path:str = '.') -> List[PythonUDTFDefinition | JavaUDTFDef
         for name, obj in inspect.getmembers(module, inspect.isclass):
             # Check if the class has the specified attribute
             if hasattr(obj, '_is_omnata_udtf'):
-                matching_classes.append(PythonUDTFDefinition(
+                matching_classes.append(UDTFDefinition(
                     name=obj._omnata_udtf_name,
+                    
                     description=obj._omnata_udtf_description,
                     params=obj._omnata_udtf_params,
                     result_columns=obj._omnata_udtf_result_columns,
                     expose_to_consumer=obj._omnata_udtf_expose_to_consumer,
                     handler=obj.__module__+'.'+obj.__name__
                 ))
-            elif isinstance(obj, JavaUDTFDefinition):
+            elif isinstance(obj, UDTFDefinition) and cast(UDTFDefinition,obj).language == 'java':
                 matching_classes.append(obj)
             
     return matching_classes
@@ -2547,7 +2492,7 @@ def omnata_udf(
     
     return decorator
 
-def find_udf_functions(path:str = '.') -> List[PythonUDFDefinition | JavaUDFDefinition]:
+def find_udf_functions(path:str = '.') -> List[UDFDefinition]:
     """
     Finds all functions in the specified directory which have the 'omnata_udf' decorator applied
     """
@@ -2566,8 +2511,11 @@ def find_udf_functions(path:str = '.') -> List[PythonUDFDefinition | JavaUDFDefi
         for name, obj in inspect.getmembers(module, inspect.isfunction):
             # Check if the class has the specified attribute
             if hasattr(obj, '_is_omnata_udf'):
-                matching_classes.append(PythonUDFDefinition(
+                matching_classes.append(UDFDefinition(
                     name=obj._omnata_udf_name,
+                    language='python',
+                    runtime_version='3.10',
+                    imports=['/app.zip'],
                     description=obj._omnata_udf_description,
                     params=obj._omnata_udf_params,
                     result_data_type=obj._omnata_udf_result_data_type,
@@ -2575,7 +2523,7 @@ def find_udf_functions(path:str = '.') -> List[PythonUDFDefinition | JavaUDFDefi
                     packages=[],
                     handler=obj.__module__+'.'+obj.__name__
                 ))
-            elif isinstance(obj, JavaUDFDefinition):
+            elif isinstance(obj, UDFDefinition) and cast(UDFDefinition,obj).language == 'java':
                 matching_classes.append(obj)
 
     return matching_classes