SNOW-4174500: Fix nested udxf inheriting inaccessable stage path

CHANGELOG.md

@@ -12,6 +12,7 @@
 - Fixed `Session.write_pandas` raising a raw connector error instead of Snowpark's table-does-not-exist message when the target table was missing and `auto_create_table=False`.
 - Fixed `ai_extract` misrouting FILE-type inputs to the TEXT overload when `scores` or `config` were also supplied.
 - Fixed the nullable `FILE` column schema expression.
+- Fixed nested UDxF registration inside a stored procedure inheriting inaccessible app-package stage paths from the procedure's `IMPORTS` clause.
 
 #### Documentation
 

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):

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)

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

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",

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