entity selector args

Author: sfc-gh-bkogan

src/snowflake/snowpark/session.py

@@ -2424,7 +2424,10 @@ def _get_default_artifact_repository(self) -> str:
             if isinstance(self._conn, MockServerConnection):
                 return _DEFAULT_ARTIFACT_REPOSITORY
 
-            cache_key = (self.get_current_database(), self.get_current_schema())
+            account = self.get_current_account()
+            database = self.get_current_database()
+            schema = self.get_current_schema()
+            cache_key = (database, schema)
 
             if (
                 self._default_artifact_repository_cache is not None
@@ -2434,8 +2437,15 @@ def _get_default_artifact_repository(self) -> str:
 
             try:
                 python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+                entity_selector_args = (
+                    f"'schema', '{schema}'"
+                    if schema
+                    else f"'database', '{database}'"
+                    if database
+                    else f"'account', '{account}'"
+                )
                 result = self._run_query(
-                    f"SELECT SYSTEM$GET_DEFAULT_PYTHON_ARTIFACT_REPOSITORY('{python_version}')"
+                    f"SELECT SYSTEM$GET_DEFAULT_PYTHON_ARTIFACT_REPOSITORY('{python_version}', {entity_selector_args})"
                 )
                 value = result[0][0] if result else None
                 resolved = value or _DEFAULT_ARTIFACT_REPOSITORY

tests/integ/test_udf.py

@@ -3108,18 +3108,17 @@ def test_urllib() -> str:
 )
 def test_use_default_artifact_repository(db_parameters):
     with Session.builder.configs(db_parameters).create() as session:
+        temp_database = Utils.random_temp_database()
         temp_schema = Utils.random_temp_schema()
-        session.sql(f"create schema {temp_schema}").collect()
-        session.sql(f"use schema {temp_schema}").collect()
+        session.sql(f"create database {temp_database}").collect()
+        session.sql(f"use database {temp_database}").collect()
         session.sql(
             "ALTER SESSION SET ENABLE_DEFAULT_PYTHON_ARTIFACT_REPOSITORY = true"
         ).collect()
         session.sql(
-            "ALTER schema set DEFAULT_PYTHON_ARTIFACT_REPOSITORY = testdb_snowpark_python.testschema_snowpark_python.SNOWPARK_PYTHON_TEST_REPOSITORY"
+            "ALTER database set DEFAULT_PYTHON_ARTIFACT_REPOSITORY = snowflake.snowpark.anaconda_shared_repository"
         ).collect()
 
-        session.add_packages("art", "cloudpickle")
-
         def test_art() -> str:
             import art  # art is not available in the conda channel, but is in pypi
 
@@ -3128,6 +3127,25 @@ def test_art() -> str:
 
         temp_func_name = Utils.random_name_for_temp_object(TempObjectType.FUNCTION)
 
+        # should not work in the database where the default is anaconda
+        with pytest.raises(
+            Exception,
+            match="Cannot add package art because it is not available in Snowflake",
+        ):
+            udf(
+                session=session,
+                func=test_art,
+                name=temp_func_name,
+                packages=["art", "cloudpickle"],
+            )
+
+        session.sql(f"create schema {temp_schema}").collect()
+        session.use_schema(temp_schema)
+        session.sql(
+            "ALTER schema set DEFAULT_PYTHON_ARTIFACT_REPOSITORY = testdb_snowpark_python.testschema_snowpark_python.SNOWPARK_PYTHON_TEST_REPOSITORY"
+        ).collect()
+        session.add_packages("art", "cloudpickle")
+
         try:
             # Test function registration
             udf(
@@ -3142,4 +3160,4 @@ def test_art() -> str:
         finally:
             session._run_query(f"drop function if exists {temp_func_name}(int)")
 
-        session.sql(f"drop schema {temp_schema}").collect()
+        session.sql(f"drop database {temp_database}").collect()