SNOW-2161990 introduce a tiny abstraction to allow sproc to override … (#2370)

Author: sfc-gh-zyao

src/snowflake/connector/file_transfer_agent.py

@@ -1062,11 +1062,14 @@ def _init_file_metadata(self) -> None:
             for idx, file_name in enumerate(self._src_files):
                 if not file_name:
                     continue
-                first_path_sep = file_name.find("/")
                 dst_file_name = (
-                    file_name[first_path_sep + 1 :]
+                    self._strip_stage_prefix_from_dst_file_name_for_download(file_name)
+                )
+                first_path_sep = dst_file_name.find("/")
+                dst_file_name = (
+                    dst_file_name[first_path_sep + 1 :]
                     if first_path_sep >= 0
-                    else file_name
+                    else dst_file_name
                 )
                 url = None
                 if self._presigned_urls and idx < len(self._presigned_urls):
@@ -1201,3 +1204,12 @@ def _process_file_compression_type(self) -> None:
                 else:
                     m.dst_file_name = m.name
                     m.dst_compression_type = None
+
+    def _strip_stage_prefix_from_dst_file_name_for_download(self, dst_file_name):
+        """Strips the stage prefix from dst_file_name for download.
+
+        Note that this is no-op in most cases, and therefore we return as is.
+        But for some workloads they will monkeypatch this method to add their
+        stripping logic.
+        """
+        return dst_file_name

test/unit/test_put_get.py

@@ -3,6 +3,7 @@
 
 from os import chmod, path
 from unittest import mock
+from unittest.mock import patch
 
 import pytest
 
@@ -252,3 +253,43 @@ def test_iobound_limit(tmp_path):
                     pass
     # 2 IObound TPEs should be created for 3 files limited to 2
     assert len(list(filter(lambda e: e.args == (2,), tpe.call_args_list))) == 2
+
+
+def test_strip_stage_prefix_from_dst_file_name_for_download():
+    """Verifies that _strip_stage_prefix_from_dst_file_name_for_download is called when initializing file meta.
+
+    Workloads like sproc will need to monkeypatch _strip_stage_prefix_from_dst_file_name_for_download on the server side
+    to maintain its behavior. So we add this unit test to make sure that we do not accidentally refactor this method and
+    break sproc workloads.
+    """
+    file = "test.txt"
+    agent = SnowflakeFileTransferAgent(
+        mock.MagicMock(autospec=SnowflakeCursor),
+        "GET @stage_foo/test.txt file:///tmp",
+        {
+            "data": {
+                "localLocation": "/tmp",
+                "command": "DOWNLOAD",
+                "autoCompress": False,
+                "src_locations": [file],
+                "sourceCompression": "none",
+                "stageInfo": {
+                    "creds": {},
+                    "location": "",
+                    "locationType": "S3",
+                    "path": "remote_loc",
+                },
+            },
+            "success": True,
+        },
+    )
+    agent._parse_command()
+    with patch.object(
+        agent,
+        "_strip_stage_prefix_from_dst_file_name_for_download",
+        return_value="mock value",
+    ):
+        agent._init_file_metadata()
+        agent._strip_stage_prefix_from_dst_file_name_for_download.assert_called_with(
+            file
+        )