diff --git a/CHANGELOG.md b/CHANGELOG.md index f44163c7c9..c317e8e93c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ #### New Features #### Bug Fixes +- Fixed nested UDxF registration inside a stored procedure inheriting inaccessible app-package stage paths from the procedure's `IMPORTS` clause. #### Improvements diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index 13b84be1d8..5cca3aa1d4 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -425,6 +425,7 @@ def __init__(self, session: "Session", conf: Dict[str, Any]) -> None: "flatten_select_after_filter_and_orderby": True, "collect_stacktrace_in_query_tag": False, "use_simplified_query_generation": False, + "restage_stored_procedure_imports": True, # SNOW-4174500 } # For config that's temporary/to be removed soon self._lock = self._session._lock for key, val in conf.items(): @@ -1601,6 +1602,59 @@ def _resolve_import_path( else: return trimmed_path, None, None + def _redirect_inherited_stage_imports( + self, + import_paths: Dict[str, Tuple[Optional[str], Optional[str]]], + ) -> Dict[str, Tuple[Optional[str], Optional[str]]]: + """Re-stage inherited session imports from the procedure's local import directory. + + Inside a stored procedure the session's import paths are seeded with the absolute + stage paths from the procedure's IMPORTS clause. Nested UDxF registrations that + inherit these paths verbatim fail with 002003/093023 when the executing role cannot + access the original stage. This method replaces each such path with its local copy + under ``snowflake_import_directory`` so it is re-uploaded to the session temp stage. + No-op outside a stored procedure or when the fix is disabled via + ``session.conf.set("restage_stored_procedure_imports", False)``. + """ + if not is_in_stored_procedure() or not self._conf.get( + "restage_stored_procedure_imports", True + ): + return import_paths + + # https://docs.snowflake.com/en/developer-guide/udf/python/udf-python-examples#reading-a-statically-specified-file-using-imports + import_directory = sys._xoptions.get("snowflake_import_directory") + if not import_directory: # pragma: no cover + return import_paths + + redirected: Dict[str, Tuple[Optional[str], Optional[str]]] = {} + for path, import_info in import_paths.items(): + local_path = None + if path.startswith(STAGE_PREFIX): + file_name = path.rsplit("/", 1)[-1].strip() + if file_name: + candidate = os.path.join(import_directory, file_name) + if os.path.isfile(candidate): + local_path = candidate + if local_path is None: + redirected[path] = import_info + continue + try: + resolved_path, checksum, leading_path = self._resolve_import_path( + local_path + ) + except Exception as ex: # pragma: no cover + _logger.warning( + "Could not re-stage inherited import %s from local copy %s: %s", + path, + local_path, + ex, + ) + redirected[path] = import_info + continue + _logger.debug("Re-staging %s via local copy %s", path, local_path) + redirected[resolved_path] = (checksum, leading_path) + return redirected + def _resolve_imports( self, import_only_stage: str, @@ -1625,7 +1679,14 @@ def _resolve_imports( ) with self._lock: + # explicit imports=[] must not trigger re-staging. + inherited_from_session = udf_level_import_paths is None import_paths = udf_level_import_paths or self._import_paths.copy() + if inherited_from_session: + # SNOW-4174500: inside a stored procedure the session holds absolute stage + # paths from the procedure's own IMPORTS clause. Re-point them at the local + # copies the runtime materializes so nested UDxFs can use them + import_paths = self._redirect_inherited_stage_imports(import_paths) for path, (prefix, leading_path) in import_paths.items(): # stage file if path.startswith(STAGE_PREFIX): diff --git a/tests/integ/test_stored_procedure.py b/tests/integ/test_stored_procedure.py index 7333571796..5de89bc8bc 100644 --- a/tests/integ/test_stored_procedure.py +++ b/tests/integ/test_stored_procedure.py @@ -8,10 +8,13 @@ import logging import os import re +import shutil +import sys as _sys import time -from typing import Dict, List, Optional, Union -from unittest.mock import patch from textwrap import dedent +from typing import Dict, Iterable, List, Optional, Tuple, Union +from unittest import mock +from unittest.mock import patch import pytest @@ -25,6 +28,7 @@ except ImportError: is_pandas_available = False +import snowflake.snowpark.session as _session_module from snowflake.snowpark import Session, AsyncJob from snowflake.snowpark.context import _DEFAULT_ARTIFACT_REPOSITORY from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted @@ -2796,3 +2800,141 @@ def promote(session: Session, m: YearMonthInterval) -> YearMonthInterval: ] # 14 months + 12 = 26 months → '+2-02' assert result == "+2-02" + + +# SNOW-4174500: nested UDxF registration must not forward IMPORTS stage paths verbatim. +def _snow4174500_setup(session, resources_path, tmp_path): + """Upload helper to a temp stage; copy locally; drop stage (makes path inaccessible). + Returns (abs_stage_path, import_dir).""" + test_files = TestFiles(resources_path) + stage = Utils.random_stage_name() + Utils.create_stage(session, stage, is_temporary=True) + Utils.upload_to_stage( + session, stage, test_files.test_pandas_apply_helper_py_file, compress=False + ) + db = session.get_current_database().strip('"') + schema = session.get_current_schema().strip('"') + abs_stage_path = f"@{db}.{schema}.{stage}/test_pandas_apply_helper.py" + shutil.copy( + test_files.test_pandas_apply_helper_py_file, + tmp_path / "test_pandas_apply_helper.py", + ) + Utils.drop_stage(session, stage) + return abs_stage_path, str(tmp_path) + os.sep + + +def _run_apply_in_pandas(session, module): + rows = ( + session.create_dataframe([(1, 5), (1, 10), (2, 20)], schema=["GRP", "V"]) + .group_by("GRP") + .apply_in_pandas( + module.double_col, + output_schema=StructType( + [StructField("GRP", IntegerType()), StructField("V", IntegerType())] + ), + ) + .collect() + ) + assert len(rows) == 3 and all(r["V"] in (10, 20, 40) for r in rows) + + +def _run_udf(session, module): + scale = session.udf.register( + module.double_scalar, return_type=IntegerType(), input_types=[IntegerType()] + ) + result = ( + session.create_dataframe([[21]], schema=["X"]) + .select(scale(col("X"))) + .collect()[0][0] + ) + assert result == 42 + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", reason="local testing mode" +) +@pytest.mark.skipif(IS_IN_STORED_PROC, reason="not meaningful inside a stored proc") +@pytest.mark.parametrize( + "run_fn,need_pandas", + [ + (_run_udf, False), + pytest.param( + _run_apply_in_pandas, + True, + marks=pytest.mark.skipif(not is_pandas_available, reason="pandas required"), + ), + ], +) +def test_nested_udxf_restages_inaccessible_inherited_import( + session, resources_path, tmp_path, run_fn, need_pandas +): + """SNOW-4174500: a UDxF registered with an inaccessible inherited stage import must + succeed. The stage is dropped so the path is genuinely unreachable; the fix re-stages + from the local copy in import_dir.""" + abs_stage_path, import_dir = _snow4174500_setup(session, resources_path, tmp_path) + _sys.path.insert(0, str(tmp_path)) + try: + import test_pandas_apply_helper # noqa: F401 — importable from tmp_path; cloudpickle by-ref + + session._import_paths[abs_stage_path] = (None, None) + try: + with mock.patch.object( + _session_module, "is_in_stored_procedure", return_value=True + ), mock.patch.dict( + _sys._xoptions, {"snowflake_import_directory": import_dir} + ): + run_fn(session, test_pandas_apply_helper) + finally: + session._import_paths.pop(abs_stage_path, None) + finally: + _sys.path.remove(str(tmp_path)) + _sys.modules.pop("test_pandas_apply_helper", None) + + +@pytest.mark.skipif( + "config.getoption('local_testing_mode', default=False)", reason="local testing mode" +) +@pytest.mark.skipif(IS_IN_STORED_PROC, reason="not meaningful inside a stored proc") +def test_nested_udtf_import_rewritten_to_session_stage_fix_vs_bug( + session, resources_path, tmp_path +): + """SNOW-4174500: fix enabled → UDTF DDL uses session temp stage; fix disabled → 002003 (bug).""" + abs_stage_path, import_dir = _snow4174500_setup(session, resources_path, tmp_path) + stage_token = abs_stage_path.split("/")[-2].split(".")[-1] + + session._import_paths[abs_stage_path] = (None, None) + try: + + class Echo: + def process(self, x: int) -> Iterable[Tuple[int]]: + yield (x,) + + output_schema = StructType([StructField("v", IntegerType())]) + + with mock.patch.object( + _session_module, "is_in_stored_procedure", return_value=True + ), mock.patch.dict(_sys._xoptions, {"snowflake_import_directory": import_dir}): + udtf_obj = session.udtf.register( + Echo, output_schema=output_schema, input_types=[IntegerType()] + ) + + ddl = session.sql( + f"SELECT GET_DDL('function', '{udtf_obj.name}(INT)')" + ).collect()[0][0] + assert stage_token.upper() not in ddl.upper() + + session.conf.set("restage_stored_procedure_imports", False) + try: + with mock.patch.object( + _session_module, "is_in_stored_procedure", return_value=True + ), mock.patch.dict( + _sys._xoptions, {"snowflake_import_directory": import_dir} + ): + with pytest.raises(Exception, match="002003|093023|does not exist"): + session.udtf.register( + Echo, output_schema=output_schema, input_types=[IntegerType()] + ) + finally: + session.conf.set("restage_stored_procedure_imports", True) + finally: + session._import_paths.pop(abs_stage_path, None) diff --git a/tests/resources/test_pandas_apply_helper.py b/tests/resources/test_pandas_apply_helper.py new file mode 100644 index 0000000000..4462a23b26 --- /dev/null +++ b/tests/resources/test_pandas_apply_helper.py @@ -0,0 +1,16 @@ +# +# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. +# +# Helper module for SNOW-4174500 integration tests. +# Must be a staged file whose module-level functions cloudpickle serializes by +# reference, triggering the nested-UDxF import-inheritance bug. + + +def double_col(pdf): + result = pdf.copy() + result["V"] = result["V"] * 2 + return result + + +def double_scalar(x): + return x * 2 diff --git a/tests/unit/scala/test_utils_suite.py b/tests/unit/scala/test_utils_suite.py index f32528370f..969f24c6bd 100644 --- a/tests/unit/scala/test_utils_suite.py +++ b/tests/unit/scala/test_utils_suite.py @@ -358,6 +358,7 @@ def check_zip_files_and_close_stream(input_stream, expected_files): "resources/testJsonSameSchema.json", "resources/test_all_data_types.parquet", "resources/test_file_with_special_characters.parquet", + "resources/test_pandas_apply_helper.py", "resources/test_requirements.txt", "resources/test_requirements_py313.txt", "resources/test_requirements_unsupported.txt", diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index a11c557ea7..a3b85f598f 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -4,6 +4,7 @@ import json import logging import os +import sys import types from typing import Optional from unittest import mock @@ -1108,3 +1109,72 @@ def test_vsc_history_exporter_not_registered_when_env_var_unset(monkeypatch): assert session._vsc_history_exporter is None fake_connection.add_query_listener.assert_not_called() + + +# SNOW-4174500: _redirect_inherited_stage_imports +_APP_PACKAGE_IMPORT = ( + '@APP_PKG_DB."APP_ARTIFACTS_V1_0.82".APP_FILES/helpers/pandas_helper.py' +) + + +def _make_import_dir(tmp_path_factory, *file_names): + d = tmp_path_factory.mktemp("import_dir") + for f in file_names: + (d / f).write_text("SCALE = 1\n") + return d + + +def test_redirect_rewrites_stage_import_inside_stored_procedure( + tmp_path_factory, mock_server_connection, monkeypatch +): + session = Session(mock_server_connection) + import_dir = _make_import_dir(tmp_path_factory, "pandas_helper.py") + local_copy = str(import_dir / "pandas_helper.py") + monkeypatch.setitem( + sys._xoptions, "snowflake_import_directory", f"{import_dir}{os.sep}" + ) + + with mock.patch.object( + snowflake.snowpark.session, "is_in_stored_procedure", return_value=True + ): + redirected = session._redirect_inherited_stage_imports( + {_APP_PACKAGE_IMPORT: (None, None)} + ) + + assert list(redirected) == [local_copy] + assert redirected[local_copy][0] # checksum set + assert redirected[local_copy][1] is None + + +@pytest.mark.parametrize( + "setup,import_paths_fn", + [ + # Not in a stored procedure — no rewrite regardless of local copy. + ("not_in_sproc", lambda d: {_APP_PACKAGE_IMPORT: (None, None)}), + # In sproc but no local copy exists for the stage file. + ("no_local_copy", lambda d: {"@other_stage/missing.py": (None, None)}), + # Local (non-stage) imports must never be rewritten. + ("local_import", lambda d: {str(d / "pandas_helper.py"): ("cksum", None)}), + # Config flag disables the fix. + ("disabled", lambda d: {_APP_PACKAGE_IMPORT: (None, None)}), + ], +) +def test_redirect_is_noop( + setup, import_paths_fn, tmp_path_factory, mock_server_connection, monkeypatch +): + session = Session(mock_server_connection) + import_dir = _make_import_dir(tmp_path_factory, "pandas_helper.py") + monkeypatch.setitem( + sys._xoptions, "snowflake_import_directory", f"{import_dir}{os.sep}" + ) + if setup == "disabled": + session.conf.set("restage_stored_procedure_imports", False) + + import_paths = import_paths_fn(import_dir) + in_sproc = setup != "not_in_sproc" + with mock.patch.object( + snowflake.snowpark.session, "is_in_stored_procedure", return_value=in_sproc + ): + result = session._redirect_inherited_stage_imports(import_paths) + + assert result == import_paths diff --git a/tests/unit/test_udtf.py b/tests/unit/test_udtf.py index b038907153..222f3c19a9 100644 --- a/tests/unit/test_udtf.py +++ b/tests/unit/test_udtf.py @@ -2,8 +2,11 @@ # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. # -from collections import defaultdict +import os as _os import sys +import sys as _sys +import threading +from collections import defaultdict from typing import Tuple from unittest import mock @@ -11,6 +14,8 @@ from snowflake.connector import ProgrammingError from snowflake.snowpark import Session +import snowflake.snowpark.session as session_module +from snowflake.snowpark._internal.udf_utils import resolve_imports_and_packages from snowflake.snowpark._internal.utils import ( TempObjectType, set_ast_state, @@ -118,3 +123,81 @@ def process(self, n: int) -> Iterable[Tuple[int]]: "schema": "some_schema", "application_roles": ["app_viewer"], } + + +# SNOW-4174500: session-imported stage paths must be re-staged, not forwarded verbatim. +_APP_STAGE_PATH = '@SAMOOHA_APP_PKG."APP_ARTIFACTS_V1_0_82".APP_FILES/pandas_helper.py' + + +def test_bug_reproduced_imports_none_inherits_all_session_paths(): + """Bug baseline: imports=None inherits all session stage paths (002003/093023 trigger).""" + session = mock.MagicMock() + session._import_paths = {_APP_STAGE_PATH: (None, None)} + session._lock = threading.RLock() + session._resolve_imports.return_value = [_APP_STAGE_PATH] + session.get_session_stage.return_value = "@TEMP_SESSION_STAGE" + session._get_default_artifact_repository.return_value = "conda_channel" + session._get_packages_by_artifact_repository.return_value = {} + session._resolve_packages.return_value = ["'cloudpickle>=3.1.1'"] + session._runtime_version_from_requirement = None + + _, _, all_imports, _, _, _ = resolve_imports_and_packages( + session=session, + object_type=TempObjectType.TABLE_FUNCTION, + func=lambda pdf: pdf, + arg_names=["pdf"], + udf_name="test_udtf", + stage_location=None, + imports=None, + packages=None, + ) + assert _APP_STAGE_PATH in all_imports + + +def test_fix_redirect_is_called_for_session_level_imports(tmp_path): + """Fix: session-level imports are re-staged; the original stage path must not appear.""" + (tmp_path / "pandas_helper.py").write_text("SCALE = 1\n") + import_dir = f"{tmp_path}{_os.sep}" + + with Session.builder.config("local_testing", True).create() as real_session: + real_session._import_paths[_APP_STAGE_PATH] = (None, None) + with mock.patch.object( + session_module, "is_in_stored_procedure", return_value=True + ), mock.patch.dict( + _sys._xoptions, {"snowflake_import_directory": import_dir} + ), mock.patch.object( + real_session, "_list_files_in_stage", return_value=set() + ), mock.patch.object( + real_session._conn, "upload_stream" + ): + resolved = real_session._resolve_imports("@TEMP_STAGE", "@TEMP_STAGE") + real_session._import_paths.pop(_APP_STAGE_PATH, None) + + assert all(_APP_STAGE_PATH not in r for r in resolved) + assert any("pandas_helper" in r for r in resolved) + + +def test_fix_explicit_udf_imports_bypass_redirect(tmp_path): + """Explicit imports= are the caller's choice and must not be redirected.""" + (tmp_path / "pandas_helper.py").write_text("SCALE = 1\n") + import_dir = f"{tmp_path}{_os.sep}" + + with Session.builder.config("local_testing", True).create() as real_session: + with mock.patch.object( + session_module, "is_in_stored_procedure", return_value=True + ), mock.patch.dict( + _sys._xoptions, {"snowflake_import_directory": import_dir} + ), mock.patch.object( + real_session, + "_redirect_inherited_stage_imports", + wraps=real_session._redirect_inherited_stage_imports, + ) as mock_redirect, mock.patch.object( + real_session, "_list_files_in_stage", return_value=set() + ), mock.patch.object( + real_session._conn, "upload_stream" + ): + real_session._resolve_imports( + "@TEMP_STAGE", "@TEMP_STAGE", {_APP_STAGE_PATH: (None, None)} + ) + + mock_redirect.assert_not_called() diff --git a/tests/utils.py b/tests/utils.py index 59f66e126e..f958941d70 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1642,6 +1642,10 @@ def test_sp_mod3_py_file(self): def test_table_sp_py_file(self): return os.path.join(self.test_sp_directory, "test_table_sp_file.py") + @property + def test_pandas_apply_helper_py_file(self): + return os.path.join(self.resources_path, "test_pandas_apply_helper.py") + @property def test_pandas_udf_py_file(self): return os.path.join(self.test_udf_directory, "test_pandas_udf_file.py")