Add new API connection for SiS/SPCS #1236

sfc-gh-smirzaei · 2024-02-02T01:35:06Z

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.

Fixes #NNNN
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.

The goal is to add a new API connect for SiS and SPCS where we can create a session using the default connection params.

sfc-gh-smirzaei · 2024-02-02T01:41:08Z

src/snowflake/snowpark/session.py

+        else:
+            try:
+                # 2) check if we are running in SPCS
+                with open("/snowflake/session/token") as file:


We could replace this block by adding a helper method like _get_default_SPSC_connection_params in connector (similar to _get_default_connection_params) to get SPCS default params and move the whole thing in _create_internal. Please lemme know which one is a better option/pattern.

I have no preference, either will work!

sfc-gh-mkeller · 2024-02-02T03:20:36Z

src/snowflake/snowpark/session.py

+        ):  # (will be removed before merging) no need for `"streamlit" in sys.modules` check since is_in_stored_procedure() must be true for SiS as well
+            # Basically noop and return the existing session
+            return _get_active_session()
+        else:


No need for the else-clause since you return in the then-clause. Just get rid of this line and de-dent the rest

sfc-gh-mkeller · 2024-02-02T03:21:56Z

src/snowflake/snowpark/session.py

+            except Exception:
+                # a general error happened, re-raise it
+                raise


No need to catch an exception just to throw it again, let it bubble up

Suggested change

except Exception:

# a general error happened, re-raise it

raise

sfc-gh-mkeller · 2024-02-02T03:22:34Z

src/snowflake/snowpark/session.py

+        else:
+            try:
+                # 2) check if we are running in SPCS
+                with open("/snowflake/session/token") as file:


I have no preference, either will work!

sfc-gh-smirzaei · 2024-02-03T00:40:10Z

src/snowflake/snowpark/session.py

+            return _get_active_session()
+        try:
+            # 2) check if we are running in SPCS
+            with open("/snowflake/session/token") as file:


Any suggestion/idea on how I can test this for SPCS?

This is probably even more challenging than stored proc. In stored proc, we have a test job that creates Anaconda package, and does some repodata hack to use the new package (test job: https://ci-dev-142.int.snowflakecomputing.com/job/PythonStoredProcBuildPrecommitTest/).

Do you know how SPCS does its regression test today?

sfc-gh-smirzaei · 2024-02-03T00:44:49Z

tests/integ/test_session.py

@@ -586,7 +585,10 @@ def test_create_session_from_default_config_file(monkeypatch, db_parameters):
        with monkeypatch.context() as m:
            m.setenv("SNOWFLAKE_CONNECTIONS", tomlkit.dumps(doc))
            m.setenv("SNOWFLAKE_DEFAULT_CONNECTION_NAME", "default")
-            with Session.builder.create() as new_session:
+            if not IS_IN_STORED_PROC:


Hmm now I'm not sure if it works (haven't tested this yet)

https://ci-dev-142.int.snowflakecomputing.com/job/PythonStoredProcBuildPrecommitTest/895/
ran a SP test for you

some tests failed. Looks like a rebase with main is needed before we do another run

sfc-gh-sfan · 2024-02-05T18:20:34Z

src/snowflake/snowpark/session.py

@@ -2749,4 +2749,32 @@ def _explain_query(self, query: str) -> Optional[str]:
            _logger.warning("query `%s` cannot be explained", query)
            return None

+    @staticmethod
+    def create() -> "Session":


Is create the right action name for "gets existing session if one exists", or should this be fetch instead?

I think there is already a method called getOrCreate

So basically this comes from a request from field/PM: https://docs.google.com/document/d/12FGgc9NeH-xa31FMZa_lO2N5Xq71Cm0aquTLI2-G0IE/edit#heading=h.hvtde6rw3c3p

They want a unified connect method in session (which I mistakenly named the method created) which works across different use-cases (i.e., when Snowpark is used as is or in SiS, SP, SPCS)

I see. @sfc-gh-mkeller @sfc-gh-yixie : Do you have concerns for introducing a connect method in session?

src/snowflake/snowpark/session.py

sfc-gh-sfan · 2024-02-05T18:21:59Z

src/snowflake/snowpark/session.py

+            return _get_active_session()
+        try:
+            # 2) check if we are running in SPCS
+            with open("/snowflake/session/token") as file:


This is probably even more challenging than stored proc. In stored proc, we have a test job that creates Anaconda package, and does some repodata hack to use the new package (test job: https://ci-dev-142.int.snowflakecomputing.com/job/PythonStoredProcBuildPrecommitTest/).

Do you know how SPCS does its regression test today?

sfc-gh-aalam · 2024-02-15T04:52:36Z

src/snowflake/snowpark/session.py

+        ):  # (will be removed before merging) no need for `"streamlit" in sys.modules` check since is_in_stored_procedure() must be true for SiS as well
+            # Basically noop and return the existing session


it would be removed before merging to main? I don't understand why have it then.

There is already Session.builder.getOrCreate()

getOrCreate as is, doesn't work (since we expect to have token in the toml file and not as a separate file). But I had a conversion with Srikanth, he is tasked with adding connect for SPCS for all connectors. He will write a design doc based on the existing connect API in Python connector where we take into account a new connect argument (like token_file) to be added to the toml file. I'll hold off working on this for now since his general solution is better/cleaner than what I have here.

sfc-gh-yixie · 2024-02-15T07:41:45Z

src/snowflake/snowpark/_internal/utils.py

@@ -501,6 +501,10 @@ def is_in_stored_procedure():
    return PLATFORM == "XP"


+def is_in_spcs():
+    return os.path.exists("/snowflake/session/token")


Is there another way?

sfc-gh-yixie · 2024-02-15T07:42:55Z

src/snowflake/snowpark/session.py

+        ):  # (will be removed before merging) no need for `"streamlit" in sys.modules` check since is_in_stored_procedure() must be true for SiS as well
+            # Basically noop and return the existing session


There is already Session.builder.getOrCreate()

sfc-gh-bhess · 2024-04-05T22:07:27Z

src/snowflake/snowpark/session.py

+                "protocol": os.getenv("SNOWFLAKE_PROTOCOL"),
+                "authenticator": "oauth",
+                "token": token,
+                "warehouse": os.getenv("SNOWFLAKE_WAREHOUSE"),


SNOWFLAKE_WAREHOUSE may or may not be set. Setting the QUERY_WAREHOUSE on the SERVICE is optional.

Add new API connection for SiS/SPCS

3b3ee01

sfc-gh-smirzaei requested a review from sfc-gh-mkeller February 2, 2024 01:35

remove redundant check

d9e44f3

sfc-gh-smirzaei commented Feb 2, 2024

View reviewed changes

sfc-gh-mkeller reviewed Feb 2, 2024

View reviewed changes

apply comments

c9158c3

sfc-gh-smirzaei requested review from sfc-gh-mkeller, sfc-gh-sfan and sfc-gh-yixie February 2, 2024 23:36

add test

d0a833a

sfc-gh-smirzaei commented Feb 3, 2024

View reviewed changes

sfc-gh-smirzaei marked this pull request as ready for review February 3, 2024 00:40

sfc-gh-smirzaei requested a review from a team as a code owner February 3, 2024 00:40

sfc-gh-smirzaei commented Feb 3, 2024

View reviewed changes

sfc-gh-sfan reviewed Feb 5, 2024

View reviewed changes

sfc-gh-smirzaei added 2 commits February 8, 2024 00:32

Applyin first round of reviews

f676f66

get rid of redundant else

12707f1

sfc-gh-smirzaei requested review from sfc-gh-aalam, sfc-gh-sfan and sfc-gh-sichen February 8, 2024 00:36

add missing args

79cf041

sfc-gh-aalam reviewed Feb 15, 2024

View reviewed changes

sfc-gh-yixie reviewed Feb 15, 2024

View reviewed changes

sfc-gh-bhess reviewed Apr 5, 2024

View reviewed changes

sfc-gh-smirzaei closed this Apr 29, 2024

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add new API connection for SiS/SPCS #1236

Add new API connection for SiS/SPCS #1236

sfc-gh-smirzaei commented Feb 2, 2024

sfc-gh-smirzaei Feb 2, 2024

sfc-gh-mkeller Feb 2, 2024

sfc-gh-mkeller Feb 2, 2024

sfc-gh-mkeller Feb 2, 2024

sfc-gh-mkeller Feb 2, 2024

sfc-gh-smirzaei Feb 3, 2024

sfc-gh-sfan Feb 5, 2024

sfc-gh-smirzaei Feb 3, 2024

sfc-gh-aalam Feb 15, 2024

sfc-gh-aalam Feb 15, 2024

sfc-gh-sfan Feb 5, 2024

sfc-gh-aalam Feb 5, 2024

sfc-gh-smirzaei Feb 8, 2024

sfc-gh-sfan Feb 8, 2024

sfc-gh-sfan Feb 5, 2024

sfc-gh-aalam Feb 15, 2024

sfc-gh-yixie Feb 15, 2024

sfc-gh-stan Feb 15, 2024

sfc-gh-smirzaei Feb 16, 2024

sfc-gh-yixie Feb 15, 2024

sfc-gh-yixie Feb 15, 2024

sfc-gh-bhess Apr 5, 2024

	except Exception:
	# a general error happened, re-raise it
	raise

		): # (will be removed before merging) no need for `"streamlit" in sys.modules` check since is_in_stored_procedure() must be true for SiS as well
		# Basically noop and return the existing session

Add new API connection for SiS/SPCS #1236

Add new API connection for SiS/SPCS #1236

Conversation

sfc-gh-smirzaei commented Feb 2, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment