SNOW-878135 is_permanent behavior fix

CHANGELOG.md

@@ -8,6 +8,10 @@
   - `array_flatten`
 - Added support for replicating your local Python environment on Snowflake via `Session.replicate_local_environment`.
 
+### Behavior Changes
+
+- When creating stored procedures, UDFs, UDTFs, UDAFs with parameter `is_permanent=False` will now create temporary objects even when `stage_name` is provided. Earlier `is_permanent=False` with non-None `stage_name` try to create permanent objects at the given stage location.
+
 ## 1.6.1 (2023-08-02)
 
 ### New Features

src/snowflake/snowpark/_internal/udf_utils.py

@@ -331,7 +331,7 @@ def check_register_args(
     is_permanent: bool = False,
     stage_location: Optional[str] = None,
     parallel: int = 4,
-):
+) -> Optional[str]:
     if is_permanent:
         if not name:
             raise ValueError(
@@ -341,12 +341,19 @@ def check_register_args(
             raise ValueError(
                 f"stage_location must be specified for permanent {get_error_message_abbr(object_type)}"
             )
+    else:
+        if stage_location:
+            logger.warn(
+                "is_permanent is False therefore stage_location will be ignored"
+            )
 
     if parallel < 1 or parallel > 99:
         raise ValueError(
             "Supported values of parallel are from 1 to 99, " f"but got {parallel}"
         )
 
+    return stage_location if is_permanent else None
+
 
 def check_execute_as_arg(execute_as: typing.Literal["caller", "owner"]):
     if (

src/snowflake/snowpark/stored_procedure.py

@@ -515,7 +515,7 @@ def register(
             )
 
         check_execute_as_arg(execute_as)
-        check_register_args(
+        stage_location = check_register_args(
             TempObjectType.PROCEDURE, name, is_permanent, stage_location, parallel
         )
 
@@ -642,7 +642,7 @@ def register_from_file(
             - :meth:`register`
         """
         file_path = process_file_path(file_path)
-        check_register_args(
+        stage_location = check_register_args(
             TempObjectType.PROCEDURE, name, is_permanent, stage_location, parallel
         )
         check_execute_as_arg(execute_as)

src/snowflake/snowpark/udaf.py

@@ -401,7 +401,7 @@ def register(
                 f"Invalid handler: expecting a class type, but get {type(handler)}"
             )
 
-        check_register_args(
+        stage_location = check_register_args(
             TempObjectType.AGGREGATE_FUNCTION,
             name,
             is_permanent,
@@ -529,7 +529,7 @@ def register_from_file(
             - :meth:`register`
         """
         file_path = process_file_path(file_path)
-        check_register_args(
+        stage_location = check_register_args(
             TempObjectType.AGGREGATE_FUNCTION,
             name,
             is_permanent,

src/snowflake/snowpark/udf.py

@@ -583,7 +583,7 @@ def register(
                 f"(__call__ is not defined): {type(func)}"
             )
 
-        check_register_args(
+        stage_location = check_register_args(
             TempObjectType.FUNCTION, name, is_permanent, stage_location, parallel
         )
 
@@ -720,7 +720,7 @@ def register_from_file(
             - :meth:`register`
         """
         file_path = process_file_path(file_path)
-        check_register_args(
+        stage_location = check_register_args(
             TempObjectType.FUNCTION, name, is_permanent, stage_location, parallel
         )
 

src/snowflake/snowpark/udtf.py

@@ -490,7 +490,7 @@ def register(
                 f"(__call__ is not defined): {type(handler)}"
             )
 
-        check_register_args(
+        stage_location = check_register_args(
             TempObjectType.TABLE_FUNCTION, name, is_permanent, stage_location, parallel
         )
 
@@ -614,7 +614,7 @@ def register_from_file(
             - :meth:`register`
         """
         file_path = process_file_path(file_path)
-        check_register_args(
+        stage_location = check_register_args(
             TempObjectType.TABLE_FUNCTION, name, is_permanent, stage_location, parallel
         )
 

tests/integ/test_stored_procedure.py

@@ -189,6 +189,7 @@ def test_call_named_stored_procedure(session, temp_schema, db_parameters):
             stage_location=unwrap_stage_location_single_quote(
                 tmp_stage_name_in_temp_schema
             ),
+            is_permanent=True,
         )
         assert new_session.call(full_sp_name, 13, 19) == 13 + 19
         # oen result in the temp schema
@@ -575,6 +576,40 @@ def test_permanent_sp(session, db_parameters):
             Utils.drop_stage(session, stage_name)
 
 
+@pytest.mark.skipif(IS_IN_STORED_PROC, reason="Cannot create session in SP")
+def test_permanent_sp_negative(session, db_parameters, caplog):
+    stage_name = Utils.random_stage_name()
+    sp_name = Utils.random_name_for_temp_object(TempObjectType.PROCEDURE)
+    with Session.builder.configs(db_parameters).create() as new_session:
+        new_session.sql_simplifier_enabled = session.sql_simplifier_enabled
+        new_session.add_packages("snowflake-snowpark-python")
+        try:
+            with caplog.at_level(logging.WARN):
+                sproc(
+                    lambda session_, x, y: session_.sql(f"SELECT {x} + {y}").collect()[
+                        0
+                    ][0],
+                    return_type=IntegerType(),
+                    input_types=[IntegerType(), IntegerType()],
+                    name=sp_name,
+                    is_permanent=False,
+                    stage_location=stage_name,
+                    session=new_session,
+                )
+            assert (
+                "is_permanent is False therefore stage_location will be ignored"
+                in caplog.text
+            )
+
+            with pytest.raises(
+                SnowparkSQLException, match=f"Unknown function {sp_name}"
+            ):
+                session.call(sp_name, 1, 2)
+            assert new_session.call(sp_name, 8, 9) == 17
+        finally:
+            new_session._run_query(f"drop function if exists {sp_name}(int, int)")
+
+
 @pytest.mark.skipif(not is_pandas_available, reason="Requires pandas")
 def test_sp_negative(session):
     def f(_, x):

tests/integ/test_udf.py

@@ -206,6 +206,7 @@ def test_call_named_udf(session, temp_schema, db_parameters):
             stage_location=unwrap_stage_location_single_quote(
                 tmp_stage_name_in_temp_schema
             ),
+            is_permanent=True,
         )
         Utils.check_answer(
             new_session.sql(f"select {full_udf_name}(13, 19)").collect(), [Row(13 + 19)]
@@ -967,6 +968,40 @@ def test_permanent_udf(session, db_parameters):
             Utils.drop_stage(session, stage_name)
 
 
+@pytest.mark.skipif(IS_IN_STORED_PROC, reason="Cannot create session in SP")
+def test_permanent_udf_negative(session, db_parameters, caplog):
+    stage_name = Utils.random_stage_name()
+    udf_name = Utils.random_name_for_temp_object(TempObjectType.FUNCTION)
+    with Session.builder.configs(db_parameters).create() as new_session:
+        new_session.sql_simplifier_enabled = session.sql_simplifier_enabled
+        try:
+            with caplog.at_level(logging.WARN):
+                udf(
+                    lambda x, y: x + y,
+                    return_type=IntegerType(),
+                    input_types=[IntegerType(), IntegerType()],
+                    name=udf_name,
+                    is_permanent=False,
+                    stage_location=stage_name,
+                    session=new_session,
+                )
+            assert (
+                "is_permanent is False therefore stage_location will be ignored"
+                in caplog.text
+            )
+
+            with pytest.raises(
+                SnowparkSQLException, match=f"Unknown function {udf_name}"
+            ):
+                session.sql(f"select {udf_name}(8, 9)").collect()
+
+            Utils.check_answer(
+                new_session.sql(f"select {udf_name}(8, 9)").collect(), [Row(17)]
+            )
+        finally:
+            new_session._run_query(f"drop function if exists {udf_name}(int, int)")
+
+
 def test_udf_negative(session):
     def f(x):
         return x