keep the PRIVATE api consistent from before

Author: sfc-gh-bkogan

src/snowflake/snowpark/_internal/udf_utils.py

@@ -1142,9 +1142,9 @@ def add_snowpark_package_to_sproc_packages(
             packages = [this_package]
         else:
             with session._package_lock:
-                existing_packages = session._artifact_repository_packages[
+                existing_packages = session._get_packages_by_artifact_repository(
                     artifact_repository
-                ]
+                )
                 if package_name not in existing_packages:
                     packages = list(existing_packages.values()) + [this_package]
         return packages
@@ -1239,7 +1239,9 @@ def resolve_imports_and_packages(
         )
 
     existing_packages_dict = (
-        session._artifact_repository_packages[artifact_repository] if session else {}
+        session._get_packages_by_artifact_repository(artifact_repository)
+        if session
+        else {}
     )
 
     if artifact_repository != _ANACONDA_SHARED_REPOSITORY:
@@ -1248,7 +1250,9 @@ def resolve_imports_and_packages(
         if not packages and session:
             resolved_packages = list(
                 session._resolve_packages(
-                    [], artifact_repository, existing_packages_dict
+                    [],
+                    artifact_repository=artifact_repository,
+                    existing_packages_dict=existing_packages_dict,
                 )
             )
         elif packages:
@@ -1286,17 +1290,17 @@ def resolve_imports_and_packages(
             resolved_packages = (
                 session._resolve_packages(
                     packages,
-                    artifact_repository,
-                    {},  # ignore session packages if passed in explicitly
+                    artifact_repository=artifact_repository,
+                    existing_packages_dict={},  # ignore session packages if passed in explicitly
                     include_pandas=is_pandas_udf,
                     statement_params=statement_params,
                     _suppress_local_package_warnings=_suppress_local_package_warnings,
                 )
                 if packages is not None
                 else session._resolve_packages(
                     [],
-                    artifact_repository,
-                    existing_packages_dict,
+                    artifact_repository=artifact_repository,
+                    existing_packages_dict=existing_packages_dict,
                     validate_package=False,
                     include_pandas=is_pandas_udf,
                     statement_params=statement_params,

src/snowflake/snowpark/session.py

@@ -601,7 +601,8 @@ def __init__(
         self._conn = conn
         self._query_tag = None
         self._import_paths: Dict[str, Tuple[Optional[str], Optional[str]]] = {}
-        # unused, needed for test infra?
+        # packages under the DEFAULT_ARTIFACT_REPOSITORY
+        # due to server side accessing private session members, this cannot be merged with _artifact_repository_packages
         self._packages: Dict[str, str] = {}
         # map of artifact repository name -> packages that should be added to functions under that repository
         self._artifact_repository_packages: DefaultDict[
@@ -1601,6 +1602,14 @@ def _list_files_in_stage(
         prefix_length = get_stage_file_prefix_length(stage_location)
         return {str(row[0])[prefix_length:] for row in file_list}
 
+    def _get_packages_by_artifact_repository(
+        self, artifact_repository: str
+    ) -> Dict[str, str]:
+        if artifact_repository == _DEFAULT_ARTIFACT_REPOSITORY:
+            return self._packages
+        else:
+            return self._artifact_repository_packages[artifact_repository]
+
     def get_packages(self, artifact_repository: Optional[str] = None) -> Dict[str, str]:
         """
         Returns a ``dict`` of packages added for user-defined functions (UDFs).
@@ -1615,7 +1624,7 @@ def get_packages(self, artifact_repository: Optional[str] = None) -> Dict[str, s
             artifact_repository = self._get_default_artifact_repository()
 
         with self._package_lock:
-            return self._artifact_repository_packages[artifact_repository].copy()
+            return self._get_packages_by_artifact_repository(artifact_repository).copy()
 
     def add_packages(
         self,
@@ -1688,8 +1697,10 @@ def add_packages(
 
         self._resolve_packages(
             parse_positional_args_to_list(*packages),
-            artifact_repository,
-            self._artifact_repository_packages[artifact_repository],
+            artifact_repository=artifact_repository,
+            existing_packages_dict=self._get_packages_by_artifact_repository(
+                artifact_repository
+            ),
         )
 
     def remove_package(
@@ -1726,7 +1737,7 @@ def remove_package(
             artifact_repository = self._get_default_artifact_repository()
 
         with self._package_lock:
-            packages = self._artifact_repository_packages[artifact_repository]
+            packages = self._get_packages_by_artifact_repository(artifact_repository)
             if package_name in packages:
                 packages.pop(package_name)
             else:
@@ -1744,7 +1755,7 @@ def clear_packages(
             artifact_repository = self._get_default_artifact_repository()
 
         with self._package_lock:
-            self._artifact_repository_packages[artifact_repository].clear()
+            self._get_packages_by_artifact_repository(artifact_repository).clear()
 
     def add_requirements(
         self,
@@ -2112,11 +2123,11 @@ def _get_req_identifiers_list(
     def _resolve_packages(
         self,
         packages: List[Union[str, ModuleType]],
-        artifact_repository: str,
-        existing_packages_dict: Dict[str, str],
+        existing_packages_dict: Dict[str, str] = None,
         validate_package: bool = True,
         include_pandas: bool = False,
         statement_params: Optional[Dict[str, str]] = None,
+        artifact_repository: str = None,
         **kwargs,
     ) -> List[str]:
         """
@@ -2134,6 +2145,13 @@ def _resolve_packages(
         Returns:
             List[str]: List of package specifiers
         """
+        if artifact_repository is None:
+            artifact_repository = self._get_default_artifact_repository()
+        if existing_packages_dict is None:
+            existing_packages_dict = self._get_packages_by_artifact_repository(
+                artifact_repository
+            )
+
         # Always include cloudpickle
         extra_modules = [cloudpickle]
         if include_pandas:

tests/integ/test_packaging.py

@@ -271,8 +271,8 @@ def extract_major_minor_patch(version_string):
 
     resolved_packages = session._resolve_packages(
         [numpy, pandas, dateutil],
-        _ANACONDA_SHARED_REPOSITORY,
-        {},
+        artifact_repository=_ANACONDA_SHARED_REPOSITORY,
+        existing_packages_dict={},
         validate_package=False,
     )
     # resolved_packages is a list of strings like
@@ -1204,17 +1204,10 @@ def test_replicate_local_environment(session):
         "force_push": True,
     }
 
-    assert not any(
-        [
-            package.startswith("cloudpickle")
-            for package in session._artifact_repository_packages[
-                _ANACONDA_SHARED_REPOSITORY
-            ]
-        ]
-    )
+    assert not any([package.startswith("cloudpickle") for package in session._packages])
 
     def naive_add_packages(self, packages):
-        self._artifact_repository_packages[_ANACONDA_SHARED_REPOSITORY] = packages
+        self._packages = packages
 
     with patch.object(session, "_is_anaconda_terms_acknowledged", lambda: True):
         with patch.object(Session, "add_packages", new=naive_add_packages):
@@ -1228,22 +1221,10 @@ def naive_add_packages(self, packages):
                 },
             )
 
-    assert any(
-        [
-            package.startswith("cloudpickle==")
-            for package in session._artifact_repository_packages[
-                _ANACONDA_SHARED_REPOSITORY
-            ]
-        ]
-    )
+    assert any([package.startswith("cloudpickle==") for package in session._packages])
     for default_package in DEFAULT_PACKAGES:
         assert not any(
-            [
-                package.startswith(default_package)
-                for package in session._artifact_repository_packages[
-                    _ANACONDA_SHARED_REPOSITORY
-                ]
-            ]
+            [package.startswith(default_package) for package in session._packages]
         )
 
     session.clear_packages()
@@ -1262,29 +1243,12 @@ def naive_add_packages(self, packages):
                 ignore_packages=ignored_packages, relax=True
             )
 
-    assert any(
-        [
-            package == "cloudpickle"
-            for package in session._artifact_repository_packages[
-                _ANACONDA_SHARED_REPOSITORY
-            ]
-        ]
-    )
+    assert any([package == "cloudpickle" for package in session._packages])
     for default_package in DEFAULT_PACKAGES:
         assert not any(
-            [
-                package.startswith(default_package)
-                for package in session._artifact_repository_packages[
-                    _ANACONDA_SHARED_REPOSITORY
-                ]
-            ]
+            [package.startswith(default_package) for package in session._packages]
         )
     for ignored_package in ignored_packages:
         assert not any(
-            [
-                package.startswith(ignored_package)
-                for package in session._artifact_repository_packages[
-                    _ANACONDA_SHARED_REPOSITORY
-                ]
-            ]
+            [package.startswith(ignored_package) for package in session._packages]
         )

tests/unit/test_session.py

@@ -213,8 +213,8 @@ def mock_get_information_schema_packages(table_name: str, _emit_ast: bool = True
 
     session._resolve_packages(
         ["random_package_name"],
-        _ANACONDA_SHARED_REPOSITORY,
-        {},
+        artifact_repository=_ANACONDA_SHARED_REPOSITORY,
+        existing_packages_dict={},
         validate_package=True,
         include_pandas=False,
     )
@@ -248,8 +248,8 @@ def run_query(sql: str):
     ):
         session._resolve_packages(
             ["random_package_name"],
-            _ANACONDA_SHARED_REPOSITORY,
-            {},
+            artifact_repository=_ANACONDA_SHARED_REPOSITORY,
+            existing_packages_dict={},
             validate_package=True,
             include_pandas=False,
         )
@@ -273,8 +273,8 @@ def mock_get_information_schema_packages(table_name: str, _emit_ast: bool = True
 
     resolved_packages = session._resolve_packages(
         ["random_package_name"],
-        _ANACONDA_SHARED_REPOSITORY,
-        existing_packages,
+        artifact_repository=_ANACONDA_SHARED_REPOSITORY,
+        existing_packages_dict=existing_packages,
         validate_package=True,
         include_pandas=False,
     )
@@ -305,8 +305,8 @@ def mock_get_information_schema_packages(table_name: str, _emit_ast: bool = True
     ):
         session._resolve_packages(
             ["snowflake-snowpark-python"],
-            _ANACONDA_SHARED_REPOSITORY,
-            {},
+            artifact_repository=_ANACONDA_SHARED_REPOSITORY,
+            existing_packages_dict={},
             validate_package=True,
             include_pandas=False,
             _suppress_local_package_warnings=True,
@@ -333,16 +333,16 @@ def assert_packages(packages):
 
     packages = session._resolve_packages(
         ["snowflake-snowpark-python==1.0.0", "cloudpickle==1.0.0"],
-        "snowflake.snowpark.pypi_shared_repository",
-        existing_packages,
+        artifact_repository="snowflake.snowpark.pypi_shared_repository",
+        existing_packages_dict=existing_packages,
     )
 
     assert_packages(packages)
 
     packages = session._resolve_packages(
         [],
-        "snowflake.snowpark.pypi_shared_repository",
-        existing_packages,
+        artifact_repository="snowflake.snowpark.pypi_shared_repository",
+        existing_packages_dict=existing_packages,
     )
 
     assert_packages(packages)

tests/unit/test_stored_procedure.py

@@ -44,6 +44,7 @@ def test_stored_procedure_execute_as(execute_as):
     fake_session._analyzer = Analyzer(fake_session)
     fake_session._runtime_version_from_requirement = None
     fake_session._artifact_repository_packages = defaultdict(dict)
+    fake_session._packages = {}
 
     def return1(_):
         return 1
@@ -92,6 +93,7 @@ def test_do_register_sp_negative(cleanup_registration_patch):
     fake_session._run_query = mock.Mock(side_effect=ProgrammingError())
     fake_session.sproc = StoredProcedureRegistration(fake_session)
     fake_session._artifact_repository_packages = defaultdict(dict)
+    fake_session._packages = {}
     with pytest.raises(SnowparkSQLException) as ex_info:
         sproc(lambda: 1, session=fake_session, return_type=IntegerType(), packages=[])
     assert ex_info.value.error_code == "1304"

tests/unit/test_udaf.py

@@ -59,6 +59,7 @@ def test_do_register_udaf_negative(cleanup_registration_patch):
     fake_session._runtime_version_from_requirement = None
     fake_session.udaf = UDAFRegistration(fake_session)
     fake_session._artifact_repository_packages = defaultdict(dict)
+    fake_session._packages = {}
     with pytest.raises(SnowparkSQLException) as ex_info:
 
         @udaf(session=fake_session)

tests/unit/test_udf.py

@@ -34,6 +34,7 @@ def test_do_register_sp_negative(cleanup_registration_patch):
     fake_session._run_query = mock.Mock(side_effect=ProgrammingError())
     fake_session.udf = UDFRegistration(fake_session)
     fake_session._artifact_repository_packages = defaultdict(dict)
+    fake_session._packages = {}
     with pytest.raises(SnowparkSQLException) as ex_info:
         udf(lambda: 1, session=fake_session, return_type=IntegerType(), packages=[])
     assert ex_info.value.error_code == "1304"

tests/unit/test_udf_utils.py

@@ -2,7 +2,6 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
-from collections import defaultdict
 import logging
 import os
 import pickle
@@ -254,8 +253,7 @@ def test_add_snowpark_package_to_sproc_packages_does_not_replace_package():
 
 def test_add_snowpark_package_to_sproc_packages_to_session():
     fake_session = mock.create_autospec(Session)
-    fake_session._artifact_repository_packages = defaultdict(dict)
-    fake_session._artifact_repository_packages[_ANACONDA_SHARED_REPOSITORY] = {
+    fake_session._packages = {
         "random_package_one": "random_package_one",
         "random_package_two": "random_package_two",
     }
@@ -272,7 +270,7 @@ def test_add_snowpark_package_to_sproc_packages_to_session():
     assert len(result) == 3
     assert final_name in result
 
-    fake_session._artifact_repository_packages[_ANACONDA_SHARED_REPOSITORY][
+    fake_session._packages[
         "snowflake-snowpark-python"
     ] = "snowflake-snowpark-python==1.12.0"
     result = add_snowpark_package_to_sproc_packages(

tests/unit/test_udtf.py

@@ -42,6 +42,7 @@ def test_do_register_sp_negative(cleanup_registration_patch):
     fake_session._runtime_version_from_requirement = None
     fake_session.udtf = UDTFRegistration(fake_session)
     fake_session._artifact_repository_packages = defaultdict(dict)
+    fake_session._packages = {}
     with pytest.raises(SnowparkSQLException) as ex_info:
 
         @udtf(output_schema=["num"], session=fake_session)