dev changes deployment for the udfs

Author: shushilgirish

scripts/deployment_files/check_and_fix_udf.py

@@ -1,6 +1,11 @@
 import os
 import sys
 import argparse
+import logging
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 def fix_udf_function(udf_path):
     """
@@ -10,7 +15,7 @@ def fix_udf_function(udf_path):
     function_file = os.path.join(udf_path, "function.py")
 
     if not os.path.exists(function_file):
-        print(f"Error: Function file not found at: {function_file}")
+        logger.error(f"Error: Function file not found at: {function_file}")
         return False
 
     # Back up original file
@@ -24,9 +29,15 @@ def fix_udf_function(udf_path):
             f.write(content)
 
         # Check if the function has a session parameter
-        if "def main(session, input_data" in content:
-            print("Found Snowpark UDF with session parameter")
+        # We need to be more thorough in our detection
+        if "def main(session" in content or "def main( session" in content:
+            logger.info("Found Snowpark UDF with session parameter")
 
+            # Check if it's already been wrapped
+            if "def main_with_session(session" in content:
+                logger.info("Function is already wrapped - no need to modify")
+                return True
+                
             # Add wrapper function that gets session from snowflake.snowpark.functions
             modified_content = content.replace(
                 "def main(session, input_data",
@@ -69,27 +80,61 @@ def main(input_data"""
             with open(function_file, 'w') as f:
                 f.write(modified_content)
 
-            print(f"✅ Updated {function_file} with wrapper function")
-            print(f"Original file backed up to {backup_file}")
+            logger.info(f"✅ Updated {function_file} with wrapper function")
+            logger.info(f"Original file backed up to {backup_file}")
             return True
         else:
-            print("No session parameter detected, no changes needed")
+            logger.info("No session parameter detected, no changes needed")
             return False
 
     except Exception as e:
-        print(f"Error fixing UDF: {e}")
+        logger.error(f"Error fixing UDF: {e}")
         # Try to restore backup if it exists
         if os.path.exists(backup_file):
-            print("Restoring backup...")
+            logger.info("Restoring backup...")
             with open(backup_file, 'r') as f:
                 original = f.read()
             with open(function_file, 'w') as f:
                 f.write(original)
         return False
 
+def analyze_udf_file(udf_path):
+    """Analyze the UDF file to determine its signature."""
+    function_file = os.path.join(udf_path, "function.py")
+    
+    if not os.path.exists(function_file):
+        logger.error(f"Function file not found: {function_file}")
+        return
+        
+    with open(function_file, 'r') as f:
+        content = f.read()
+    
+    logger.info(f"Analyzing UDF in {function_file}")
+    
+    # Extract the main function signature
+    import re
+    sig_match = re.search(r'def\s+main\s*\((.*?)\)', content)
+    if sig_match:
+        params = sig_match.group(1)
+        logger.info(f"Function signature parameters: '{params}'")
+        if "session" in params:
+            logger.info("This UDF uses the Snowpark session parameter and needs wrapping")
+            return True
+    
+    return False
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Fix UDF function for Snowflake compatibility')
     parser.add_argument('udf_path', help='Path to UDF directory containing function.py')
+    parser.add_argument('--analyze', action='store_true', help='Just analyze, don\'t fix')
 
     args = parser.parse_args()
-    fix_udf_function(args.udf_path)
+    
+    if args.analyze:
+        analyze_udf_file(args.udf_path)
+    else:
+        needs_fix = analyze_udf_file(args.udf_path)
+        if needs_fix:
+            fix_udf_function(args.udf_path)
+        else:
+            logger.info("UDF doesn't need fixing")

scripts/deployment_files/snowflake_deployer.py

@@ -147,6 +147,41 @@ def create_snowflake_connection(conn_config):
         logger.error(f"Failed to create Snowflake connection: {str(e)}")
         raise
 
+def analyze_function_signature(function_file):
+    """Analyze the function signature to determine parameter structure."""
+    try:
+        import re
+        with open(function_file, 'r') as f:
+            content = f.read()
+        
+        # Look for the main function definition
+        signature_match = re.search(r'def\s+main\s*\((.*?)\)', content)
+        if signature_match:
+            params = signature_match.group(1).strip()
+            logger.info(f"Function signature parameters: '{params}'")
+            
+            # Count parameters (excluding session if present)
+            param_list = [p.strip() for p in params.split(',') if p.strip()]
+            
+            # Check if session is a parameter
+            has_session = any(p.strip().startswith('session') for p in param_list)
+            param_count = len(param_list)
+            
+            return {
+                'has_session': has_session,
+                'param_count': param_count,
+                'param_list': param_list
+            }
+    except Exception as e:
+        logger.error(f"Error analyzing function signature: {e}")
+    
+    # Default fallback
+    return {
+        'has_session': False,
+        'param_count': 1,
+        'param_list': ['input_data']
+    }
+
 def deploy_component(profile_name, component_path, component_name, component_type):
     """Deploy a component to Snowflake."""
     logger.info(f"Deploying component: {component_name} ({component_type})")
@@ -196,9 +231,14 @@ def deploy_component(profile_name, component_path, component_name, component_typ
                     logger.info(f"Using actual code directory: {code_dir}")
 
             # Check and fix UDF function signature if necessary
-            if component_type.lower() == "udf":
-                logger.info(f"Checking and fixing UDF function signature for {component_name}")
-                os.system(f"python scripts/deployment_files/check_and_fix_udf.py {code_dir}")
+            function_file = os.path.join(code_dir, "function.py")
+            signature_info = {'has_session': False, 'param_count': 1, 'param_list': ['input_data']}
+            
+            if component_type.lower() == "udf" and os.path.exists(function_file):
+                logger.info(f"Analyzing UDF function signature for {component_name}")
+                signature_info = analyze_function_signature(function_file)
+                logger.info(f"Function analysis: Session={signature_info['has_session']}, "
+                           f"Param Count={signature_info['param_count']}")
 
             # Log directory contents
             logger.info(f"Component directory structure:")
@@ -224,16 +264,39 @@ def deploy_component(profile_name, component_path, component_name, component_typ
             import_path = f"@{stage_name}/{component_name.replace(' ', '_')}/{zip_filename}"
 
             if component_type.lower() == "udf":
-                # For Snowpark UDFs - different SQL based on parameter signature
-                sql = f"""
-                CREATE OR REPLACE FUNCTION {component_name.replace(' ', '_')}(input_data VARIANT)
-                RETURNS VARIANT
-                LANGUAGE PYTHON
-                RUNTIME_VERSION=3.8
-                PACKAGES = ('snowflake-snowpark-python')
-                HANDLER = 'function.main'
-                IMPORTS = ('{import_path}')
-                """
+                # Adjust SQL based on parameter count
+                if signature_info['param_count'] == 1:
+                    sql = f"""
+                    CREATE OR REPLACE FUNCTION {component_name.replace(' ', '_')}(input_data VARIANT)
+                    RETURNS VARIANT
+                    LANGUAGE PYTHON
+                    RUNTIME_VERSION=3.8
+                    PACKAGES = ('snowflake-snowpark-python')
+                    IMPORTS = ('{import_path}')
+                    HANDLER = 'function.main'
+                    """
+                elif signature_info['param_count'] == 2:
+                    # For two parameters
+                    sql = f"""
+                    CREATE OR REPLACE FUNCTION {component_name.replace(' ', '_')}(current_value FLOAT, previous_value FLOAT)
+                    RETURNS FLOAT
+                    LANGUAGE PYTHON
+                    RUNTIME_VERSION=3.8
+                    PACKAGES = ('snowflake-snowpark-python')
+                    IMPORTS = ('{import_path}')
+                    HANDLER = 'function.main'
+                    """
+                else:
+                    # Fallback to standard variant parameter
+                    sql = f"""
+                    CREATE OR REPLACE FUNCTION {component_name.replace(' ', '_')}(input_data VARIANT)
+                    RETURNS VARIANT
+                    LANGUAGE PYTHON
+                    RUNTIME_VERSION=3.8
+                    PACKAGES = ('snowflake-snowpark-python')
+                    IMPORTS = ('{import_path}')
+                    HANDLER = 'function.main'
+                    """
             else:  # procedure
                 sql = f"""
                 CREATE OR REPLACE PROCEDURE {component_name.replace(' ', '_')}()