SNOW-2014274: Add restricted caller permission for execute_as (#3226)

Author: sfc-gh-jdu

CHANGELOG.md

@@ -4,9 +4,14 @@
 
 ### Snowpark Python API Updates
 
+#### New Features
+
+- Added support for `restricted caller` permission of `execute_as` argument in `StoredProcedure.regsiter()`
+
 #### Bug Fixes
 
 - Fixed a bug in `DataFrame.group_by().pivot().agg` when the pivot column and aggregate column are the same.
+- Fixed a bug in local testing where transient `__pycache__` directory was unintentionally copied during stored procedure execution via import.
 
 #### Deprecations
 

src/snowflake/snowpark/_internal/ast/utils.py

@@ -1442,7 +1442,7 @@ def build_sproc(  # type: ignore[no-untyped-def] # TODO(SNOW-1491199) # Function
     secrets: Optional[Dict[str, str]] = None,
     comment: Optional[str] = None,
     statement_params: Optional[Dict[str, str]] = None,
-    execute_as: typing.Literal["caller", "owner"] = "owner",
+    execute_as: typing.Literal["caller", "owner", "restricted caller"] = "owner",
     source_code_display: bool = True,
     is_permanent: bool = False,
     session: "snowflake.snowpark.session.Session" = None,

src/snowflake/snowpark/_internal/udf_utils.py

@@ -92,7 +92,7 @@
 AGGREGATE_FUNCTION_MERGE_METHOD = "merge"
 AGGREGATE_FUNCTION_STATE_METHOD = "aggregate_state"
 
-EXECUTE_AS_WHITELIST = frozenset(["owner", "caller"])
+EXECUTE_AS_WHITELIST = frozenset(["owner", "caller", "restricted caller"])
 
 REGISTER_KWARGS_ALLOWLIST = {
     "native_app_params",
@@ -142,7 +142,9 @@ def __init__(
         func: Union[Callable, Tuple[str, str]],
         replace: bool = False,
         if_not_exists: bool = False,
-        execute_as: Optional[typing.Literal["caller", "owner"]] = None,
+        execute_as: Optional[
+            typing.Literal["caller", "owner", "restricted caller"]
+        ] = None,
         anonymous: bool = False,
     ) -> None:
         self.func = func
@@ -513,7 +515,9 @@ def check_register_args(
         )
 
 
-def check_execute_as_arg(execute_as: typing.Literal["caller", "owner"]):
+def check_execute_as_arg(
+    execute_as: typing.Literal["caller", "owner", "restricted caller"]
+):
     if (
         not isinstance(execute_as, str)
         or execute_as.lower() not in EXECUTE_AS_WHITELIST
@@ -1266,7 +1270,7 @@ def create_python_udf_or_sp(
     if_not_exists: bool,
     raw_imports: Optional[List[Union[str, Tuple[str, str]]]],
     inline_python_code: Optional[str] = None,
-    execute_as: Optional[typing.Literal["caller", "owner"]] = None,
+    execute_as: Optional[typing.Literal["caller", "owner", "restricted caller"]] = None,
     api_call_source: Optional[str] = None,
     strict: bool = False,
     secure: bool = False,

src/snowflake/snowpark/functions.py

@@ -10448,7 +10448,7 @@ def sproc(
     session: Optional["snowflake.snowpark.Session"] = None,
     parallel: int = 4,
     statement_params: Optional[Dict[str, str]] = None,
-    execute_as: typing.Literal["caller", "owner"] = "owner",
+    execute_as: typing.Literal["caller", "owner", "restricted caller"] = "owner",
     strict: bool = False,
     source_code_display: bool = True,
     external_access_integrations: Optional[List[str]] = None,

src/snowflake/snowpark/mock/_stored_procedure.py

@@ -48,7 +48,7 @@ def __init__(
         input_types: List[DataType],
         name: str,
         imports: Set[str],
-        execute_as: typing.Literal["caller", "owner"] = "owner",
+        execute_as: typing.Literal["caller", "owner", "restricted caller"] = "owner",
         anonymous_sp_sql: Optional[str] = None,
         strict=False,
         _ast: Optional[proto.Expr] = None,
@@ -249,7 +249,7 @@ def _do_register_sp(
         *,
         source_code_display: bool = False,
         statement_params: Optional[Dict[str, str]] = None,
-        execute_as: typing.Literal["caller", "owner"] = "owner",
+        execute_as: typing.Literal["caller", "owner", "restricted caller"] = "owner",
         anonymous: bool = False,
         api_call_source: str,
         skip_upload_on_content_match: bool = False,

src/snowflake/snowpark/mock/_util.py

@@ -319,7 +319,10 @@ def cleanup_imports():
                     sys.path.append(module_path)
                 if os.path.isdir(module_path):
                     shutil.copytree(
-                        module_path, temporary_import_path.name, dirs_exist_ok=True
+                        module_path,
+                        temporary_import_path.name,
+                        dirs_exist_ok=True,
+                        ignore=shutil.ignore_patterns("__pycache__"),
                     )
                 else:
                     shutil.copy2(module_path, temporary_import_path.name)

src/snowflake/snowpark/stored_procedure.py

@@ -75,7 +75,7 @@ def __init__(
         return_type: DataType,
         input_types: List[DataType],
         name: str,
-        execute_as: typing.Literal["caller", "owner"] = "owner",
+        execute_as: typing.Literal["caller", "owner", "restricted caller"] = "owner",
         anonymous_sp_sql: Optional[str] = None,
         packages: Optional[List[Union[str, ModuleType]]] = None,
         _ast: Optional[proto.StoredProcedure] = None,
@@ -506,7 +506,7 @@ def register(
         replace: bool = False,
         if_not_exists: bool = False,
         parallel: int = 4,
-        execute_as: typing.Literal["caller", "owner"] = "owner",
+        execute_as: typing.Literal["caller", "owner", "restricted caller"] = "owner",
         strict: bool = False,
         external_access_integrations: Optional[List[str]] = None,
         secrets: Optional[Dict[str, str]] = None,
@@ -654,7 +654,7 @@ def register_from_file(
         replace: bool = False,
         if_not_exists: bool = False,
         parallel: int = 4,
-        execute_as: typing.Literal["caller", "owner"] = "owner",
+        execute_as: typing.Literal["caller", "owner", "restricted caller"] = "owner",
         strict: bool = False,
         external_access_integrations: Optional[List[str]] = None,
         secrets: Optional[Dict[str, str]] = None,
@@ -725,7 +725,7 @@ def register_from_file(
                 Increasing the number of threads can improve performance when uploading
                 large stored procedure files.
             execute_as: What permissions should the procedure have while executing. This
-                supports caller, or owner for now.
+                supports ``caller``, ``owner`` or ``restricted caller`` for now.
             strict: Whether the created stored procedure is strict. A strict stored procedure will not invoke
                 the stored procedure if any input is null. Instead, a null value will always be returned. Note
                 that the stored procedure might still return null for non-null inputs.
@@ -811,7 +811,7 @@ def _do_register_sp(
         *,
         source_code_display: bool = False,
         statement_params: Optional[Dict[str, str]] = None,
-        execute_as: typing.Literal["caller", "owner"] = "owner",
+        execute_as: typing.Literal["caller", "owner", "restricted caller"] = "owner",
         anonymous: bool = False,
         api_call_source: str,
         skip_upload_on_content_match: bool = False,

tests/integ/test_stored_procedure.py

@@ -1699,7 +1699,7 @@ def plus1(_: Session, x: int) -> int:
         session._run_query(f"drop procedure if exists {perm_sp_name}(int)")
 
 
-@pytest.mark.parametrize("execute_as", [None, "owner", "caller"])
+@pytest.mark.parametrize("execute_as", [None, "owner", "caller", "restricted caller"])
 def test_execute_as_options(session, execute_as):
     """Make sure that a stored procedure can be run with any EXECUTE AS option."""
 
@@ -1716,7 +1716,7 @@ def return1(_):
     assert return1_sp() == 1
 
 
-@pytest.mark.parametrize("execute_as", [None, "owner", "caller"])
+@pytest.mark.parametrize("execute_as", [None, "owner", "caller", "restricted caller"])
 def test_execute_as_options_while_registering_from_file(
     session, resources_path, tmpdir, execute_as
 ):