SNOW-2436917: fix udxf/sproc registration regression (#3927)

Author: sfc-gh-aling

src/snowflake/snowpark/_internal/udf_utils.py

@@ -1134,7 +1134,7 @@ def resolve_imports_and_packages(
     skip_upload_on_content_match: bool = False,
     is_permanent: bool = False,
     force_inline_code: bool = False,
-    **kwargs,
+    _suppress_local_package_warnings: bool = False,
 ) -> Tuple[
     Optional[str],
     Optional[str],
@@ -1168,9 +1168,7 @@ def resolve_imports_and_packages(
                     packages,
                     include_pandas=is_pandas_udf,
                     statement_params=statement_params,
-                    _suppress_local_package_warnings=kwargs.get(
-                        "_suppress_local_package_warnings", False
-                    ),
+                    _suppress_local_package_warnings=_suppress_local_package_warnings,
                 )
                 if packages is not None
                 else session._resolve_packages(
@@ -1179,9 +1177,7 @@ def resolve_imports_and_packages(
                     validate_package=False,
                     include_pandas=is_pandas_udf,
                     statement_params=statement_params,
-                    _suppress_local_package_warnings=kwargs.get(
-                        "_suppress_local_package_warnings", False
-                    ),
+                    _suppress_local_package_warnings=_suppress_local_package_warnings,
                 )
             )
 

src/snowflake/snowpark/stored_procedure.py

@@ -955,7 +955,10 @@ def _do_register_sp(
             is_permanent=is_permanent,
             force_inline_code=force_inline_code,
             artifact_repository=artifact_repository,
-            **kwargs,
+            _suppress_local_package_warnings=kwargs.get(
+                "_suppress_local_package_warnings", False
+            ),
+            # DO NOT pass **kwargs here, as it can lead to TypeError: multiple values for the same argument
         )
 
         runtime_version_from_requirement = None

src/snowflake/snowpark/udaf.py

@@ -799,7 +799,10 @@ def _do_register_udaf(
             skip_upload_on_content_match=skip_upload_on_content_match,
             is_permanent=is_permanent,
             artifact_repository=artifact_repository,
-            **kwargs,
+            _suppress_local_package_warnings=kwargs.get(
+                "_suppress_local_package_warnings", False
+            ),
+            # DO NOT pass **kwargs here, as it can lead to TypeError: multiple values for the same argument
         )
 
         runtime_version_from_requirement = None

src/snowflake/snowpark/udf.py

@@ -994,7 +994,10 @@ def _do_register_udf(
             skip_upload_on_content_match=skip_upload_on_content_match,
             is_permanent=is_permanent,
             artifact_repository=artifact_repository,
-            **kwargs,
+            _suppress_local_package_warnings=kwargs.get(
+                "_suppress_local_package_warnings", False
+            ),
+            # DO NOT pass **kwargs here, as it can lead to TypeError: multiple values for the same argument
         )
 
         runtime_version_from_requirement = None

src/snowflake/snowpark/udtf.py

@@ -1056,7 +1056,10 @@ def _do_register_udtf(
             skip_upload_on_content_match=skip_upload_on_content_match,
             is_permanent=is_permanent,
             artifact_repository=artifact_repository,
-            **kwargs,
+            _suppress_local_package_warnings=kwargs.get(
+                "_suppress_local_package_warnings", False
+            )
+            # DO NOT pass **kwargs here, as it can lead to TypeError: multiple values for the same argument
         )
 
         runtime_version_from_requirement = None

tests/integ/test_stored_procedure.py

@@ -561,6 +561,18 @@ def test_session_register_sp(session, local_testing_mode):
     )
     assert add_sp(1, 2) == 3
 
+    # testing SNOW-2436917
+    add_sp_passing_session = session.sproc.register(
+        lambda session_, x, y: session_.create_dataframe([(x, y)])
+        .to_df("a", "b")
+        .select(col("a") + col("b"))
+        .collect()[0][0],
+        session=session,
+        return_type=IntegerType(),
+        input_types=[IntegerType(), IntegerType()],
+    )
+    assert add_sp_passing_session(1, 2) == 3
+
     query_tag = f"QUERY_TAG_{Utils.random_alphanumeric_str(10)}"
     add_sp = session.sproc.register(
         lambda session_, x, y: session_.create_dataframe([(x, y)])

tests/integ/test_udaf.py

@@ -63,6 +63,36 @@ def finish(self):
     Utils.assert_executed_with_query_tag(session, query_tag)
 
 
+def test_basic_udaf_snow_2436917(session):
+    class PythonCountUDAFHandler:
+        def __init__(self) -> None:
+            self._count = 0
+
+        @property
+        def aggregate_state(self):
+            return self._count
+
+        def accumulate(self, input_value):
+            if input_value is not None:
+                self._count += 1
+
+        def merge(self, other_count):
+            self._count += other_count
+
+        def finish(self):
+            return self._count
+
+    count_udaf = session.udaf.register(
+        PythonCountUDAFHandler,
+        return_type=IntegerType(),
+        input_types=[IntegerType()],
+        immutable=True,
+        session=session,
+    )
+    df = session.create_dataframe([[1], [2], [None], [3]]).to_df("a")
+    Utils.check_answer(df.agg(count_udaf("a")), [Row(3)])
+
+
 # TODO: use data class as state. This triggers a bug in UDF server during pickling/unpickling of a state.
 def test_int(session):
     @udaf

tests/integ/test_udf.py

@@ -396,6 +396,23 @@ def test_session_register_udf(session, local_testing_mode):
             Row(7),
         ],
     )
+
+    # testing SNOW-SNOW-2436917
+    add_udf_passing_session = session.udf.register(
+        lambda x, y: x + y,
+        session=session,
+        return_type=IntegerType(),
+        input_types=[IntegerType(), IntegerType()],
+    )
+    assert isinstance(add_udf_passing_session.func, Callable)
+    Utils.check_answer(
+        df.select(add_udf_passing_session("a", "b")).collect(),
+        [
+            Row(3),
+            Row(7),
+        ],
+    )
+
     # Query tags not supported in local testing.
     if not local_testing_mode:
         query_tag = f"QUERY_TAG_{Utils.random_alphanumeric_str(10)}"

tests/integ/test_udtf.py

@@ -1335,6 +1335,17 @@ def process(
     ddl_sql = f"select get_ddl('FUNCTION', '{echo_udtf.name}(number)')"
     assert comment in session.sql(ddl_sql).collect()[0][0]
 
+    # testing SNOW-SNOW-2436917
+    echo_udtf_passing_session = session.udtf.register(
+        EchoUDTF,
+        output_schema=["num"],
+        comment=comment,
+        session=session,
+    )
+
+    ddl_sql = f"select get_ddl('FUNCTION', '{echo_udtf_passing_session.name}(number)')"
+    assert comment in session.sql(ddl_sql).collect()[0][0]
+
 
 @pytest.mark.parametrize("from_file", [True, False])
 @pytest.mark.parametrize(