[Async] Apply #2303 to async code

# Conflicts:
#	src/snowflake/connector/aio/_direct_file_operation_utils.py

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_build_upload_agent.py renamed to src/snowflake/connector/aio/_bind_upload_agent.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import os
 from io import BytesIO
 from logging import getLogger
 from typing import TYPE_CHECKING, cast
@@ -56,8 +57,11 @@ async def upload(self) -> None:
                 if row_idx >= len(self.rows) or size >= self._stream_buffer_size:
                     break
             try:
-                await self.cursor.execute(
-                    f"PUT file://{row_idx}.csv {self.stage_path}", file_stream=f
+                f.seek(0)
+                await self.cursor._upload_stream(
+                    input_stream=f,
+                    stage_location=os.path.join(self.stage_path, f"{row_idx}.csv"),
+                    options={"source_compression": "auto_detect"},
                 )
             except Error as err:
                 logger.debug("Failed to upload the bindings file to stage.")

src/snowflake/connector/aio/_cursor.py

@@ -23,7 +23,7 @@
     ProgrammingError,
 )
 from snowflake.connector._sql_util import get_file_transfer_type
-from snowflake.connector.aio._build_upload_agent import BindUploadAgent
+from snowflake.connector.aio._bind_upload_agent import BindUploadAgent
 from snowflake.connector.aio._result_batch import (
     ResultBatch,
     create_batches_from_response,
@@ -803,7 +803,7 @@ async def executemany(
                 bind_stage = None
                 if (
                     bind_size
-                    > self.connection._session_parameters[
+                    >= self.connection._session_parameters[
                         "CLIENT_STAGE_ARRAY_BINDING_THRESHOLD"
                     ]
                     > 0

src/snowflake/connector/aio/_direct_file_operation_utils.py

@@ -1,7 +1,15 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ._connection import SnowflakeConnection
+
+import os
 from abc import ABC, abstractmethod
 
+from ..constants import CMD_TYPE_UPLOAD
+
 
 class FileOperationParserBase(ABC):
     """The interface of internal utility functions for file operation parsing."""
@@ -37,8 +45,8 @@ async def download_as_stream(self, ret, decompress=False):
 
 
 class FileOperationParser(FileOperationParserBase):
-    def __init__(self, connection):
-        pass
+    def __init__(self, connection: SnowflakeConnection):
+        self._connection = connection
 
     async def parse_file_operation(
         self,
@@ -49,7 +57,27 @@ async def parse_file_operation(
         options,
         has_source_from_stream=False,
     ):
-        raise NotImplementedError("parse_file_operation is not yet supported")
+        """Parses a file operation by constructing SQL and getting the SQL parsing result from server."""
+        options = options or {}
+        options_in_sql = " ".join(f"{k}={v}" for k, v in options.items())
+
+        if command_type == CMD_TYPE_UPLOAD:
+            if has_source_from_stream:
+                stage_location, unprefixed_local_file_name = os.path.split(
+                    stage_location
+                )
+                local_file_name = "file://" + unprefixed_local_file_name
+            sql = f"PUT {local_file_name} ? {options_in_sql}"
+            params = [stage_location]
+        else:
+            raise NotImplementedError(f"unsupported command type: {command_type}")
+
+        async with self._connection.cursor() as cursor:
+            # Send constructed SQL to server and get back parsing result.
+            processed_params = cursor._connection._process_params_qmarks(params, cursor)
+            return await cursor._execute_helper(
+                sql, binding_params=processed_params, is_internal=True
+            )
 
 
 class StreamDownloader(StreamDownloaderBase):

test/integ/aio/test_direct_file_operation_utils_async.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import os
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, AsyncGenerator, Callable, Coroutine
+
+import pytest
+
+try:
+    from snowflake.connector.options import pandas
+    from snowflake.connector.pandas_tools import (
+        _iceberg_config_statement_helper,
+        write_pandas,
+    )
+except ImportError:
+    pandas = None
+    write_pandas = None
+    _iceberg_config_statement_helper = None
+
+if TYPE_CHECKING:
+    from snowflake.connector.aio import SnowflakeConnection, SnowflakeCursor
+
+from ..test_direct_file_operation_utils import _normalize_windows_local_path
+
+
+async def _validate_upload_content(
+    expected_content, cursor, stage_name, local_dir, base_file_name, is_compressed
+):
+    gz_suffix = ".gz"
+    stage_path = f"@{stage_name}/{base_file_name}"
+    local_path = os.path.join(local_dir, base_file_name)
+
+    await cursor.execute(
+        f"GET {stage_path} 'file://{_normalize_windows_local_path(local_dir)}'",
+    )
+    if is_compressed:
+        stage_path += gz_suffix
+        local_path += gz_suffix
+        import gzip
+
+        with gzip.open(local_path, "r") as f:
+            read_content = f.read().decode("utf-8")
+            assert read_content == expected_content, (read_content, expected_content)
+    else:
+        with open(local_path) as f:
+            read_content = f.read()
+    assert read_content == expected_content, (read_content, expected_content)
+
+
+async def _test_runner(
+    conn_cnx: Callable[..., AsyncGenerator[SnowflakeConnection]],
+    task: Callable[[SnowflakeCursor, str, str, str], Coroutine[None, None, None]],
+    is_compressed: bool,
+    special_stage_name: str = None,
+    special_base_file_name: str = None,
+):
+    from snowflake.connector._utils import TempObjectType, random_name_for_temp_object
+
+    async with conn_cnx() as conn:
+        cursor = conn.cursor()
+        stage_name = special_stage_name or random_name_for_temp_object(
+            TempObjectType.STAGE
+        )
+        await cursor.execute(f"CREATE OR REPLACE SCOPED TEMP STAGE {stage_name}")
+        expected_content = "hello, world"
+        with TemporaryDirectory() as temp_dir:
+            base_file_name = special_base_file_name or "test.txt"
+            src_file_name = os.path.join(temp_dir, base_file_name)
+            with open(src_file_name, "w") as f:
+                f.write(expected_content)
+            # Run the file operation
+            await task(cursor, stage_name, temp_dir, base_file_name)
+            # Clean up before validation.
+            os.remove(src_file_name)
+            # Validate result.
+            await _validate_upload_content(
+                expected_content,
+                cursor,
+                stage_name,
+                temp_dir,
+                base_file_name,
+                is_compressed=is_compressed,
+            )
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.parametrize("is_compressed", [False, True])
+async def test_upload(
+    conn_cnx: Callable[..., AsyncGenerator[SnowflakeConnection]],
+    is_compressed: bool,
+):
+    async def upload_task(cursor, stage_name, temp_dir, base_file_name):
+        await cursor._upload(
+            local_file_name=f"'file://{_normalize_windows_local_path(os.path.join(temp_dir, base_file_name))}'",
+            stage_location=f"@{stage_name}",
+            options={"auto_compress": is_compressed},
+        )
+
+    await _test_runner(conn_cnx, upload_task, is_compressed=is_compressed)
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.parametrize("is_compressed", [False, True])
+async def test_upload_stream(
+    conn_cnx: Callable[..., AsyncGenerator[SnowflakeConnection]],
+    is_compressed: bool,
+):
+    async def upload_stream_task(cursor, stage_name, temp_dir, base_file_name):
+        with open(f"{os.path.join(temp_dir, base_file_name)}", "rb") as input_stream:
+            await cursor._upload_stream(
+                input_stream=input_stream,
+                stage_location=f"@{os.path.join(stage_name, base_file_name)}",
+                options={"auto_compress": is_compressed},
+            )
+
+    await _test_runner(conn_cnx, upload_stream_task, is_compressed=is_compressed)

test/unit/aio/test_bind_upload_agent_async.py

@@ -9,20 +9,22 @@
 
 
 async def test_bind_upload_agent_uploading_multiple_files():
-    from snowflake.connector.aio._build_upload_agent import BindUploadAgent
+    from snowflake.connector.aio._bind_upload_agent import BindUploadAgent
 
     csr = AsyncMock(auto_spec=True)
     rows = [bytes(10)] * 10
     agent = BindUploadAgent(csr, rows, stream_buffer_size=10)
     await agent.upload()
-    assert csr.execute.call_count == 11  # 1 for stage creation + 10 files
+    assert csr.execute.call_count == 1  # 1 for stage creation
+    assert csr._upload_stream.call_count == 10  # 10 for 10 files
 
 
 async def test_bind_upload_agent_row_size_exceed_buffer_size():
-    from snowflake.connector.aio._build_upload_agent import BindUploadAgent
+    from snowflake.connector.aio._bind_upload_agent import BindUploadAgent
 
     csr = AsyncMock(auto_spec=True)
     rows = [bytes(15)] * 10
     agent = BindUploadAgent(csr, rows, stream_buffer_size=10)
     await agent.upload()
-    assert csr.execute.call_count == 11  # 1 for stage creation + 10 files
+    assert csr.execute.call_count == 1  # 1 for stage creation
+    assert csr._upload_stream.call_count == 10  # 10 for 10 files