SNOW-1963078 Port _upload / _download / _upload_stream / _download_st… (#2198)

(cherry picked from commit a3229c3)

Author: sfc-gh-zyao

src/snowflake/connector/connection.py

@@ -82,6 +82,7 @@
     PYTHON_VERSION,
     SNOWFLAKE_CONNECTOR_VERSION,
 )
+from .direct_file_operation_utils import FileOperationParser, StreamDownloader
 from .errorcode import (
     ER_CONNECTION_IS_CLOSED,
     ER_FAILED_PROCESSING_PYFORMAT,
@@ -512,6 +513,10 @@ def __init__(
         # check SNOW-1218851 for long term improvement plan to refactor ocsp code
         atexit.register(self._close_at_exit)
 
+        # Set up the file operation parser and stream downloader.
+        self._file_operation_parser = FileOperationParser(self)
+        self._stream_downloader = StreamDownloader(self)
+
     # Deprecated
     @property
     def insecure_mode(self) -> bool:

src/snowflake/connector/cursor.py

@@ -42,6 +42,8 @@
 from ._utils import _TrackedQueryCancellationTimer
 from .bind_upload_agent import BindUploadAgent, BindUploadError
 from .constants import (
+    CMD_TYPE_DOWNLOAD,
+    CMD_TYPE_UPLOAD,
     FIELD_NAME_TO_ID,
     PARAMETER_PYTHON_CONNECTOR_QUERY_RESULT_FORMAT,
     FileTransferType,
@@ -1740,6 +1742,151 @@ def get_result_batches(self) -> list[ResultBatch] | None:
         )
         return self._result_set.batches
 
+    def _download(
+        self,
+        stage_location: str,
+        target_directory: str,
+        options: dict[str, Any],
+        _do_reset: bool = True,
+    ) -> None:
+        """Downloads from the stage location to the target directory.
+
+        Args:
+            stage_location (str): The location of the stage to download from.
+            target_directory (str): The destination directory to download into.
+            options (dict[str, Any]): The download options.
+            _do_reset (bool, optional): Whether to reset the cursor before
+                downloading, by default we will reset the cursor.
+        """
+        from .file_transfer_agent import SnowflakeFileTransferAgent
+
+        if _do_reset:
+            self.reset()
+
+        # Interpret the file operation.
+        ret = self.connection._file_operation_parser.parse_file_operation(
+            stage_location=stage_location,
+            local_file_name=None,
+            target_directory=target_directory,
+            command_type=CMD_TYPE_DOWNLOAD,
+            options=options,
+        )
+
+        # Execute the file operation based on the interpretation above.
+        file_transfer_agent = SnowflakeFileTransferAgent(
+            self,
+            "",  # empty command because it is triggered by directly calling this util not by a SQL query
+            ret,
+        )
+        file_transfer_agent.execute()
+        self._init_result_and_meta(file_transfer_agent.result())
+
+    def _upload(
+        self,
+        local_file_name: str,
+        stage_location: str,
+        options: dict[str, Any],
+        _do_reset: bool = True,
+    ) -> None:
+        """Uploads the local file to the stage location.
+
+        Args:
+            local_file_name (str): The local file to be uploaded.
+            stage_location (str): The stage location to upload the local file to.
+            options (dict[str, Any]): The upload options.
+            _do_reset (bool, optional): Whether to reset the cursor before
+                uploading, by default we will reset the cursor.
+        """
+        from .file_transfer_agent import SnowflakeFileTransferAgent
+
+        if _do_reset:
+            self.reset()
+
+        # Interpret the file operation.
+        ret = self.connection._file_operation_parser.parse_file_operation(
+            stage_location=stage_location,
+            local_file_name=local_file_name,
+            target_directory=None,
+            command_type=CMD_TYPE_UPLOAD,
+            options=options,
+        )
+
+        # Execute the file operation based on the interpretation above.
+        file_transfer_agent = SnowflakeFileTransferAgent(
+            self,
+            "",  # empty command because it is triggered by directly calling this util not by a SQL query
+            ret,
+        )
+        file_transfer_agent.execute()
+        self._init_result_and_meta(file_transfer_agent.result())
+
+    def _download_stream(
+        self, stage_location: str, decompress: bool = False
+    ) -> IO[bytes]:
+        """Downloads from the stage location as a stream.
+
+        Args:
+            stage_location (str): The location of the stage to download from.
+            decompress (bool, optional): Whether to decompress the file, by
+                default we do not decompress.
+
+        Returns:
+            IO[bytes]: A stream to read from.
+        """
+        # Interpret the file operation.
+        ret = self.connection._file_operation_parser.parse_file_operation(
+            stage_location=stage_location,
+            local_file_name=None,
+            target_directory=None,
+            command_type=CMD_TYPE_DOWNLOAD,
+            options=None,
+            has_source_from_stream=True,
+        )
+
+        # Set up stream downloading based on the interpretation and return the stream for reading.
+        return self.connection._stream_downloader.download_as_stream(ret, decompress)
+
+    def _upload_stream(
+        self,
+        input_stream: IO[bytes],
+        stage_location: str,
+        options: dict[str, Any],
+        _do_reset: bool = True,
+    ) -> None:
+        """Uploads content in the input stream to the stage location.
+
+        Args:
+            input_stream (IO[bytes]): A stream to read from.
+            stage_location (str): The location of the stage to upload to.
+            options (dict[str, Any]): The upload options.
+            _do_reset (bool, optional): Whether to reset the cursor before
+                uploading, by default we will reset the cursor.
+        """
+        from .file_transfer_agent import SnowflakeFileTransferAgent
+
+        if _do_reset:
+            self.reset()
+
+        # Interpret the file operation.
+        ret = self.connection._file_operation_parser.parse_file_operation(
+            stage_location=stage_location,
+            local_file_name=None,
+            target_directory=None,
+            command_type=CMD_TYPE_UPLOAD,
+            options=options,
+            has_source_from_stream=input_stream,
+        )
+
+        # Execute the file operation based on the interpretation above.
+        file_transfer_agent = SnowflakeFileTransferAgent(
+            self,
+            "",  # empty command because it is triggered by directly calling this util not by a SQL query
+            ret,
+            source_from_stream=input_stream,
+        )
+        file_transfer_agent.execute()
+        self._init_result_and_meta(file_transfer_agent.result())
+
 
 class DictCursor(SnowflakeCursor):
     """Cursor returning results in a dictionary."""

src/snowflake/connector/direct_file_operation_utils.py

@@ -0,0 +1,64 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+
+class FileOperationParserBase(ABC):
+    """The interface of internal utility functions for file operation parsing."""
+
+    @abstractmethod
+    def __init__(self, connection):
+        pass
+
+    @abstractmethod
+    def parse_file_operation(
+        self,
+        stage_location,
+        local_file_name,
+        target_directory,
+        command_type,
+        options,
+        has_source_from_stream=False,
+    ):
+        """Converts the file operation details into a SQL and returns the SQL parsing result."""
+        pass
+
+
+class StreamDownloaderBase(ABC):
+    """The interface of internal utility functions for stream downloading of file."""
+
+    @abstractmethod
+    def __init__(self, connection):
+        pass
+
+    @abstractmethod
+    def download_as_stream(self, ret, decompress=False):
+        pass
+
+
+class FileOperationParser(FileOperationParserBase):
+    def __init__(self, connection):
+        pass
+
+    def parse_file_operation(
+        self,
+        stage_location,
+        local_file_name,
+        target_directory,
+        command_type,
+        options,
+        has_source_from_stream=False,
+    ):
+        raise NotImplementedError("parse_file_operation is not yet supported")
+
+
+class StreamDownloader(StreamDownloaderBase):
+    def __init__(self, connection):
+        pass
+
+    def download_as_stream(self, ret, decompress=False):
+        raise NotImplementedError("download_as_stream is not yet supported")

test/integ/test_connection.py

@@ -1628,3 +1628,12 @@ def test_no_auth_connection_negative_case():
     # connection is not able to run any query
     with pytest.raises(DatabaseError, match="Connection is closed"):
         conn.execute_string("select 1")
+
+
+# _file_operation_parser and _stream_downloader are newly introduced and
+# therefore should not be tested on old drivers.
+@pytest.mark.skipolddriver
+def test_file_utils_sanity_check():
+    conn = create_connection("default")
+    assert hasattr(conn._file_operation_parser, "parse_file_operation")
+    assert hasattr(conn._stream_downloader, "download_as_stream")

test/unit/test_cursor.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import time
+from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -99,3 +100,96 @@ def mock_cmd_query(*args, **kwargs):
 
     # query cancel request should be sent upon timeout
     assert mockCancelQuery.called
+
+
+# The _upload/_download/_upload_stream/_download_stream are newly introduced
+# and therefore should not be tested in old drivers.
+@pytest.mark.skipolddriver
+class TestUploadDownloadMethods(TestCase):
+    """Test the _upload/_download/_upload_stream/_download_stream methods."""
+
+    @patch("snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent")
+    def test_download(self, MockFileTransferAgent):
+        cursor, fake_conn, mock_file_transfer_agent_instance = self._setup_mocks(
+            MockFileTransferAgent
+        )
+
+        # Call _download method
+        cursor._download("@st", "/tmp/test.txt", {})
+
+        # In the process of _download execution, we expect these methods to be called
+        #   - parse_file_operation in connection._file_operation_parser
+        #   - execute in SnowflakeFileTransferAgent
+        # And we do not expect this method to be involved
+        #   - download_as_stream of connection._stream_downloader
+        fake_conn._file_operation_parser.parse_file_operation.assert_called_once()
+        fake_conn._stream_downloader.download_as_stream.assert_not_called()
+        mock_file_transfer_agent_instance.execute.assert_called_once()
+
+    @patch("snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent")
+    def test_upload(self, MockFileTransferAgent):
+        cursor, fake_conn, mock_file_transfer_agent_instance = self._setup_mocks(
+            MockFileTransferAgent
+        )
+
+        # Call _upload method
+        cursor._upload("/tmp/test.txt", "@st", {})
+
+        # In the process of _upload execution, we expect these methods to be called
+        #   - parse_file_operation in connection._file_operation_parser
+        #   - execute in SnowflakeFileTransferAgent
+        # And we do not expect this method to be involved
+        #   - download_as_stream of connection._stream_downloader
+        fake_conn._file_operation_parser.parse_file_operation.assert_called_once()
+        fake_conn._stream_downloader.download_as_stream.assert_not_called()
+        mock_file_transfer_agent_instance.execute.assert_called_once()
+
+    @patch("snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent")
+    def test_download_stream(self, MockFileTransferAgent):
+        cursor, fake_conn, mock_file_transfer_agent_instance = self._setup_mocks(
+            MockFileTransferAgent
+        )
+
+        # Call _download_stream method
+        cursor._download_stream("@st/test.txt", decompress=True)
+
+        # In the process of _download_stream execution, we expect these methods to be called
+        #   - parse_file_operation in connection._file_operation_parser
+        #   - download_as_stream of connection._stream_downloader
+        # And we do not expect this method to be involved
+        #   - execute in SnowflakeFileTransferAgent
+        fake_conn._file_operation_parser.parse_file_operation.assert_called_once()
+        fake_conn._stream_downloader.download_as_stream.assert_called_once()
+        mock_file_transfer_agent_instance.execute.assert_not_called()
+
+    @patch("snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent")
+    def test_upload_stream(self, MockFileTransferAgent):
+        cursor, fake_conn, mock_file_transfer_agent_instance = self._setup_mocks(
+            MockFileTransferAgent
+        )
+
+        # Call _upload_stream method
+        fd = MagicMock()
+        cursor._upload_stream(fd, "@st/test.txt", {})
+
+        # In the process of _upload_stream execution, we expect these methods to be called
+        #   - parse_file_operation in connection._file_operation_parser
+        #   - execute in SnowflakeFileTransferAgent
+        # And we do not expect this method to be involved
+        #   - download_as_stream of connection._stream_downloader
+        fake_conn._file_operation_parser.parse_file_operation.assert_called_once()
+        fake_conn._stream_downloader.download_as_stream.assert_not_called()
+        mock_file_transfer_agent_instance.execute.assert_called_once()
+
+    def _setup_mocks(self, MockFileTransferAgent):
+        mock_file_transfer_agent_instance = MockFileTransferAgent.return_value
+        mock_file_transfer_agent_instance.execute.return_value = None
+
+        fake_conn = FakeConnection()
+        fake_conn._file_operation_parser = MagicMock()
+        fake_conn._stream_downloader = MagicMock()
+
+        cursor = SnowflakeCursor(fake_conn)
+        cursor.reset = MagicMock()
+        cursor._init_result_and_meta = MagicMock()
+        return cursor, fake_conn, mock_file_transfer_agent_instance