SNOW-1654536: async binding stage bind upload agent (#2069)

Author: sfc-gh-aling

src/snowflake/connector/aio/_build_upload_agent.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+from io import BytesIO
+from logging import getLogger
+from typing import TYPE_CHECKING, cast
+
+from snowflake.connector import Error
+from snowflake.connector.bind_upload_agent import BindUploadAgent as BindUploadAgentSync
+from snowflake.connector.errors import BindUploadError
+
+if TYPE_CHECKING:
+    from snowflake.connector.aio import SnowflakeCursor
+
+logger = getLogger(__name__)
+
+
+class BindUploadAgent(BindUploadAgentSync):
+    def __init__(
+        self,
+        cursor: SnowflakeCursor,
+        rows: list[bytes],
+        stream_buffer_size: int = 1024 * 1024 * 10,
+    ) -> None:
+        super().__init__(cursor, rows, stream_buffer_size)
+        self.cursor = cast("SnowflakeCursor", cursor)
+
+    async def _create_stage(self) -> None:
+        await self.cursor.execute(self._CREATE_STAGE_STMT)
+
+    async def upload(self) -> None:
+        try:
+            await self._create_stage()
+        except Error as err:
+            self.cursor.connection._session_parameters[
+                "CLIENT_STAGE_ARRAY_BINDING_THRESHOLD"
+            ] = 0
+            logger.debug("Failed to create stage for binding.")
+            raise BindUploadError from err
+
+        row_idx = 0
+        while row_idx < len(self.rows):
+            f = BytesIO()
+            size = 0
+            while True:
+                f.write(self.rows[row_idx])
+                size += len(self.rows[row_idx])
+                row_idx += 1
+                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
+                )
+            except Error as err:
+                logger.debug("Failed to upload the bindings file to stage.")
+                raise BindUploadError from err
+            f.close()

src/snowflake/connector/aio/_cursor.py

@@ -27,6 +27,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._result_batch import (
     ResultBatch,
     create_batches_from_response,
@@ -746,9 +747,10 @@ async def executemany(
                 ):
                     # bind stage optimization
                     try:
-                        raise NotImplementedError(
-                            "Bind stage is not supported yet in async."
-                        )
+                        rows = self.connection._write_params_to_byte_rows(seqparams)
+                        bind_uploader = BindUploadAgent(self, rows)
+                        await bind_uploader.upload()
+                        bind_stage = bind_uploader.stage_path
                     except BindUploadError:
                         logger.debug(
                             "Failed to upload binds to stage, sending binds to "

test/helpers.py

@@ -21,6 +21,7 @@
 from snowflake.connector.compat import OK
 
 if TYPE_CHECKING:
+    import snowflake.connector.aio
     import snowflake.connector.connection
 
 try: