Allow UDF/Sproc Decorators to execute in a Local Sandbox #1341

sfc-gh-bgoel · 2024-03-28T13:18:05Z

Please answer these questions before submitting your pull requests. Thanks!

What GitHub issue is this PR addressing? Make sure that there is an accompanying issue to your PR.

There is no open Github issue for this PR as it is initiated internally within Snowflake.
Fill out the following pre-review checklist:
- I am adding a new automated test(s) to verify correctness of my new code
- I am adding new logging messages
- I am adding a new telemetry message
- I am adding new credentials
- I am adding a new dependency
Please describe how your code solves the related issue.

Background

When a user wants to create a UDF/UDTF/UDAF (UDF for short) or an Sproc, they are able to use the corresponding decorator and provide certain parameters that help in the creation of the said object in their Snowflake account. As it is, the current code takes in those parameters, performs some validation and transformation on them, pulls in some relevant session information, and then creates a SQL query that is sent to be run in the Snowflake account.

Ask

As Snowpark is a critical component of developing a Snowflake Native App, we would like to leverage this existing code so that a user does not have to duplicate their effort in writing Snowpark-only code and Snowpark-related code for their Snowflake Native App. We want to capture any user code in python files that use the above decorators, and be able to use that information outside the Snowpark environment without talking to Snowflake.

Design

In the previously-discussed consensus, Snowflake Native App, or really any user who wants to inspect the information before creating the object in Snowflake, will install a callback into the Snowpark execution, and it is through this callback that the information will be sent back to the user for inspection before creating the object.
Further, Snowflake Native App (or a user that just wants to run code in a local sandbox instead of connecting to Snowflake) logic needs to execute in a sandbox because the UDFs/SPROCs should not be immediately created in Snowflake. Snowflake Native App wants to be able to use the resolved properties to create our own SQL to be run as part of our application code, and this workflow should not happen in the account that the user is developing for. Hence the need to have a sandbox instead of the actual account. We do not want any unintended side-effects happening in the user's account and therefore want to wall-off any interaction with Snowflake.
Additionally, because we want to give control to the caller environment (an external user/client) to decide whether they should allow creation of the object in Snowflake. If they choose to not proceed with the creation, then the code should be able to receive that decision and short circuit before it can interact with Snowflake.

github-actions · 2024-03-28T13:18:19Z

CLA Assistant Lite bot All contributors have signed the CLA ✍️ ✅

src/snowflake/snowpark/context.py

src/snowflake/snowpark/_internal/udf_utils.py

src/snowflake/snowpark/context.py

src/snowflake/snowpark/_internal/udf_utils.py

src/snowflake/snowpark/functions.py

src/snowflake/snowpark/udf.py

src/snowflake/snowpark/context.py

src/snowflake/snowpark/functions.py

src/snowflake/snowpark/session.py

sfc-gh-bgoel · 2024-04-03T15:21:39Z

I have read the CLA Document and I hereby sign the CLA

sfc-gh-bgoel · 2024-04-03T15:22:28Z

recheck

sfc-gh-bdufour · 2024-04-03T20:35:33Z

src/snowflake/snowpark/functions.py

@@ -7047,6 +7047,9 @@ def udf(
    external_access_integrations: Optional[List[str]] = None,
    secrets: Optional[Dict[str, str]] = None,
    immutable: bool = False,
+    native_app_params: Optional[


Discussion point: nesting the native app-specific params this way. The thought here was that native app params don't take effect in Snowpark, and while we'd update the docs to explain it, developers don't often read the docs and are likely to be surprised. Isolating the parameters with no effect in Snowpark seemed possibly worthwhile.

sfc-gh-bdufour · 2024-04-03T20:37:51Z

src/snowflake/snowpark/mock/_connection.py

+        elif (
+            not self._local_testing
+        ):  # i.e. we are in a sandbox environment = not a true testing environment
+            return {"data": [("Statement executed successfully.",)], "sfqid": None}


Why do we need this?

During resolve_imports_and_packages, there is a call made to session.get_session_stage. This method tries to create a stage in the current session if one does not exist. This further calls connection.run_query(). If the connection is of type MockServerConnection, then any query other than use ... fails. This is because they have the mock object names, so use ... can technically be mimic-ed. But all others throw an error, and I want to prevent the error in sandbox.

I debated short circuiting before the call to session.get_session_stage at all, but figured to keep the change as pushed further down as possible. Open to changing this though.

src/snowflake/snowpark/session.py

src/snowflake/snowpark/_internal/udf_utils.py

…ce decision to register objects with Snowflake

…fy context.py to decouple sandboxing and interrupting registration

… and interrupt registration logic, short circuited some package and import level methods

sfc-gh-jrose

Overall I think the change makes sense. Separating out the native apps params seems like the right thing to do.

My main concern for this change is that it adds additional complexity to already lengthy and complex functions like resolve_imports_and_packages. That function should probably be broken up a bit into some helper functions. I usually follow the rule of thumb that if the whole function can't fit on screen at once it's probably too long.

Adding complexity to support a different execution mode to the Session class will make code maintenance harder. Instead I'd rather see a subclass of Session with appropriate parts changed. The session builder can then choose which type of session is more applicable.

Finally, global variables for control flow could probably be better handled with an environment variable for _is_execution_environment_sandboxed and a session builder parameter for _should_continue_registration. I think we should avoid mutable global variables when possible.

Those are my thoughts, but I'd be interested in seeing what other people say before you start refactoring.

sfc-gh-bdufour · 2024-04-05T18:08:49Z

Thanks for your feedback! Answers inline

Overall I think the change makes sense. Separating out the native apps params seems like the right thing to do.

My main concern for this change is that it adds additional complexity to already lengthy and complex functions like resolve_imports_and_packages. That function should probably be broken up a bit into some helper functions. I usually follow the rule of thumb that if the whole function can't fit on screen at once it's probably too long.

Fair point, however we were asked very specifically to keep the number of lines changed to a minimum and avoid refactoring for this change at least. Yijun thought that introducing refactorings would increase the risk associated with the PR, and that it was best to avoid this.

Adding complexity to support a different execution mode to the Session class will make code maintenance harder. Instead I'd rather see a subclass of Session with appropriate parts changed. The session builder can then choose which type of session is more applicable.

Keep in mind what we're doing is truly global and not scoped to a session, since we're loading arbitrary user code in a sandbox. We don't know how the user code will create sessions (e.g. through a builder or not).

Finally, global variables for control flow could probably be better handled with an environment variable for _is_execution_environment_sandboxed and a session builder parameter for _should_continue_registration. I think we should avoid mutable global variables when possible.

That was deliberate, see above regarding session state. It's actually one of those rare cases where mutable global state is desirable, IHMO. We need to guarantee that no matter how the user's code is written, we're get the information we need and prevent accidental SF queries.

…tests working on integ/test_udf add more integ tests

sfc-gh-yixie · 2024-04-10T16:23:43Z

src/snowflake/snowpark/_internal/udf_utils.py

@@ -869,7 +947,7 @@ def resolve_imports_and_packages(
                    "or a tuple of the file path (str) and the import path (str)."
                )
            udf_level_imports[resolved_import_tuple[0]] = resolved_import_tuple[1:]
-        all_urls = session._resolve_imports(
+        all_urls = session._resolve_imports(  # _resolve_imports will return an empty list in case of sandbox


For offline discussion, is it possible to skip using session if the native apps sandbox doesn't have a session?

sfc-gh-bgoel · 2024-04-16T19:09:16Z

Closing in favor of #1381

sfc-gh-bdufour reviewed Mar 28, 2024

View reviewed changes

sfc-gh-bgoel mentioned this pull request Apr 1, 2024

[Draft][Do Not Merge] Refactoring UDF Annotation Code #1332

Closed

5 tasks

sfc-gh-bdufour reviewed Apr 2, 2024

View reviewed changes

sfc-gh-bgoel requested review from sfc-gh-yixie and sfc-gh-jrose April 3, 2024 15:14

sfc-gh-bgoel closed this Apr 3, 2024

sfc-gh-bgoel reopened this Apr 3, 2024

github-actions bot locked and limited conversation to collaborators Apr 3, 2024

sfc-gh-bgoel changed the title ~~[Draft][Do Not Merge] Allow UDF Decorator in a Local Sandbox~~ Allow UDF/Sproc Decorators to execute in a Local Sandbox Apr 3, 2024

sfc-gh-bdufour reviewed Apr 3, 2024

View reviewed changes

src/snowflake/snowpark/_internal/udf_utils.py Outdated Show resolved Hide resolved

sfc-gh-bdufour reviewed Apr 3, 2024

View reviewed changes

src/snowflake/snowpark/_internal/udf_utils.py Outdated Show resolved Hide resolved

sfc-gh-stan self-requested a review April 4, 2024 21:52

sfc-gh-bgoel added 8 commits April 5, 2024 09:00

initial commit

4035572

Add comments for the flow, include schema and application roles info

6e8d1d1

conditionally resolve imports and packages if in native apps, introdu…

45ef1ff

…ce decision to register objects with Snowflake

distinction between sandbox and regular when do not register

85a1939

remove package and imports as keys from native app params, and simpli…

73d3354

…fy context.py to decouple sandboxing and interrupting registration

introduced local_testing flag for mock connection, simplified sandbox…

e791ed1

… and interrupt registration logic, short circuited some package and import level methods

addressing feedback

aad0ab8

add unit tests round 1. but tests/mock tests failing on main branch

2ad9570

sfc-gh-bgoel force-pushed the bgoel-na-hook-3-no-refactor branch from 8764dd4 to 2ad9570 Compare April 5, 2024 13:21

sfc-gh-jrose reviewed Apr 5, 2024

View reviewed changes

unit tests round 2, int tests starting

765523a

sfc-gh-bgoel marked this pull request as ready for review April 5, 2024 19:09

sfc-gh-bgoel requested a review from a team as a code owner April 5, 2024 19:09

sfc-gh-bgoel requested a review from sfc-gh-jrose April 5, 2024 19:09

sfc-gh-bgoel added 3 commits April 8, 2024 11:32

unit/test_session untouched, integ/test_session amended with session …

ba4ac32

…tests working on integ/test_udf add more integ tests

some minor changes

5fc8be4

Try with real connection, no tests, changing public api as well

37393ae

sfc-gh-yixie reviewed Apr 10, 2024

View reviewed changes

sfc-gh-bgoel closed this Apr 16, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow UDF/Sproc Decorators to execute in a Local Sandbox #1341

Allow UDF/Sproc Decorators to execute in a Local Sandbox #1341

sfc-gh-bgoel commented Mar 28, 2024 •

edited

Loading

github-actions bot commented Mar 28, 2024 •

edited

Loading

sfc-gh-bgoel commented Apr 3, 2024 •

edited

Loading

sfc-gh-bgoel commented Apr 3, 2024

sfc-gh-bdufour Apr 3, 2024 •

edited

Loading

sfc-gh-bdufour Apr 3, 2024

sfc-gh-bgoel Apr 4, 2024

sfc-gh-jrose left a comment

sfc-gh-bdufour commented Apr 5, 2024

sfc-gh-yixie Apr 10, 2024

sfc-gh-bgoel commented Apr 16, 2024

Allow UDF/Sproc Decorators to execute in a Local Sandbox #1341

Allow UDF/Sproc Decorators to execute in a Local Sandbox #1341

Conversation

sfc-gh-bgoel commented Mar 28, 2024 • edited Loading

Background

Ask

Design

github-actions bot commented Mar 28, 2024 • edited Loading

sfc-gh-bgoel commented Apr 3, 2024 • edited Loading

sfc-gh-bgoel commented Apr 3, 2024

sfc-gh-bdufour Apr 3, 2024 • edited Loading

Choose a reason for hiding this comment

sfc-gh-bdufour Apr 3, 2024

Choose a reason for hiding this comment

sfc-gh-bgoel Apr 4, 2024

Choose a reason for hiding this comment

sfc-gh-jrose left a comment

Choose a reason for hiding this comment

sfc-gh-bdufour commented Apr 5, 2024

sfc-gh-yixie Apr 10, 2024

Choose a reason for hiding this comment

sfc-gh-bgoel commented Apr 16, 2024

sfc-gh-bgoel commented Mar 28, 2024 •

edited

Loading

github-actions bot commented Mar 28, 2024 •

edited

Loading

sfc-gh-bgoel commented Apr 3, 2024 •

edited

Loading

sfc-gh-bdufour Apr 3, 2024 •

edited

Loading