snowflake-connector-python 4.4.0__cp314-cp314-win_amd64.whl

snowflake/connector/aio/_description.py

@@ -0,0 +1,5 @@
+"""Various constants."""
+
+from __future__ import annotations
+
+CLIENT_NAME = "AsyncioPythonConnector"  # don't change!

snowflake/connector/aio/_direct_file_operation_utils.py

@@ -0,0 +1,88 @@
+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."""
+
+    @abstractmethod
+    def __init__(self, connection):
+        pass
+
+    @abstractmethod
+    async 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
+    async def download_as_stream(self, ret, decompress=False):
+        pass
+
+
+class FileOperationParser(FileOperationParserBase):
+    def __init__(self, connection: SnowflakeConnection):
+        self._connection = connection
+
+    async def parse_file_operation(
+        self,
+        stage_location,
+        local_file_name,
+        target_directory,
+        command_type,
+        options,
+        has_source_from_stream=False,
+    ):
+        """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):
+    def __init__(self, connection):
+        pass
+
+    async def download_as_stream(self, ret, decompress=False):
+        raise NotImplementedError("download_as_stream is not yet supported")

snowflake/connector/aio/_file_transfer_agent.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import sys
+from logging import getLogger
+from typing import IO, TYPE_CHECKING, Any
+
+from ..constants import (
+    AZURE_CHUNK_SIZE,
+    AZURE_FS,
+    CMD_TYPE_DOWNLOAD,
+    CMD_TYPE_UPLOAD,
+    GCS_FS,
+    LOCAL_FS,
+    S3_FS,
+    ResultStatus,
+    megabyte,
+)
+from ..errorcode import ER_FILE_NOT_EXISTS
+from ..errors import Error, OperationalError
+from ..file_transfer_agent import SnowflakeFileMeta
+from ..file_transfer_agent import (
+    SnowflakeFileTransferAgent as SnowflakeFileTransferAgentSync,
+)
+from ..file_transfer_agent import SnowflakeProgressPercentage, _chunk_size_calculator
+from ..local_storage_client import SnowflakeLocalStorageClient
+from ._azure_storage_client import SnowflakeAzureRestClient
+from ._gcs_storage_client import SnowflakeGCSRestClient
+from ._s3_storage_client import SnowflakeS3RestClient
+from ._storage_client import SnowflakeStorageClient
+
+if TYPE_CHECKING:  # pragma: no cover
+    from ._cursor import SnowflakeCursor
+
+
+logger = getLogger(__name__)
+
+
+class SnowflakeFileTransferAgent(SnowflakeFileTransferAgentSync):
+    """Snowflake File Transfer Agent provides cloud provider independent implementation for putting/getting files."""
+
+    def __init__(
+        self,
+        cursor: SnowflakeCursor,
+        command: str,
+        ret: dict[str, Any],
+        put_callback: type[SnowflakeProgressPercentage] | None = None,
+        put_azure_callback: type[SnowflakeProgressPercentage] | None = None,
+        put_callback_output_stream: IO[str] = sys.stdout,
+        get_callback: type[SnowflakeProgressPercentage] | None = None,
+        get_azure_callback: type[SnowflakeProgressPercentage] | None = None,
+        get_callback_output_stream: IO[str] = sys.stdout,
+        show_progress_bar: bool = True,
+        raise_put_get_error: bool = True,
+        force_put_overwrite: bool = True,
+        skip_upload_on_content_match: bool = False,
+        multipart_threshold: int | None = None,
+        source_from_stream: IO[bytes] | None = None,
+        use_s3_regional_url: bool = False,
+        unsafe_file_write: bool = False,
+        reraise_error_in_file_transfer_work_function: bool = False,
+    ) -> None:
+        super().__init__(
+            cursor=cursor,
+            command=command,
+            ret=ret,
+            put_callback=put_callback,
+            put_azure_callback=put_azure_callback,
+            put_callback_output_stream=put_callback_output_stream,
+            get_callback=get_callback,
+            get_azure_callback=get_azure_callback,
+            get_callback_output_stream=get_callback_output_stream,
+            show_progress_bar=show_progress_bar,
+            raise_put_get_error=raise_put_get_error,
+            force_put_overwrite=force_put_overwrite,
+            skip_upload_on_content_match=skip_upload_on_content_match,
+            multipart_threshold=multipart_threshold,
+            source_from_stream=source_from_stream,
+            use_s3_regional_url=use_s3_regional_url,
+            unsafe_file_write=unsafe_file_write,
+            reraise_error_in_file_transfer_work_function=reraise_error_in_file_transfer_work_function,
+        )
+
+    async def execute(self) -> None:
+        self._parse_command()
+        self._init_file_metadata()
+
+        if self._command_type == CMD_TYPE_UPLOAD:
+            self._process_file_compression_type()
+
+        for m in self._file_metadata:
+            m.sfagent = self
+
+        await self._transfer_accelerate_config()
+
+        if self._command_type == CMD_TYPE_DOWNLOAD:
+            if not os.path.isdir(self._local_location):
+                os.makedirs(self._local_location)
+
+        if self._stage_location_type == LOCAL_FS:
+            if not os.path.isdir(self._stage_info["location"]):
+                os.makedirs(self._stage_info["location"])
+
+        for m in self._file_metadata:
+            m.overwrite = self._overwrite
+            m.skip_upload_on_content_match = self._skip_upload_on_content_match
+            m.sfagent = self
+            if self._stage_location_type != LOCAL_FS:
+                m.put_callback = self._put_callback
+                m.put_azure_callback = self._put_azure_callback
+                m.put_callback_output_stream = self._put_callback_output_stream
+                m.get_callback = self._get_callback
+                m.get_azure_callback = self._get_azure_callback
+                m.get_callback_output_stream = self._get_callback_output_stream
+                m.show_progress_bar = self._show_progress_bar
+
+                # multichunk threshold
+                m.multipart_threshold = self._multipart_threshold
+
+        # TODO: SNOW-1625364 for renaming client_prefetch_threads in asyncio
+        logger.debug(f"parallel=[{self._parallel}]")
+        if self._raise_put_get_error and not self._file_metadata:
+            Error.errorhandler_wrapper(
+                self._cursor.connection,
+                self._cursor,
+                OperationalError,
+                {
+                    "msg": "While getting file(s) there was an error: "
+                    "the file does not exist.",
+                    "errno": ER_FILE_NOT_EXISTS,
+                },
+            )
+        await self.transfer(self._file_metadata)
+
+        # turn enum to string, in order to have backward compatible interface
+
+        for result in self._results:
+            result.result_status = result.result_status.value
+
+    async def transfer(self, metas: list[SnowflakeFileMeta]) -> None:
+        files = [await self._create_file_transfer_client(m) for m in metas]
+        is_upload = self._command_type == CMD_TYPE_UPLOAD
+        finish_download_upload_tasks = []
+
+        async def preprocess_done_cb(
+            success: bool,
+            result: Any,
+            done_client: SnowflakeStorageClient,
+        ) -> None:
+            if not success:
+                logger.debug(f"Failed to prepare {done_client.meta.name}.")
+                try:
+                    if is_upload:
+                        await done_client.finish_upload()
+                        done_client.delete_client_data()
+                    else:
+                        await done_client.finish_download()
+                except Exception as error:
+                    done_client.meta.error_details = error
+            elif done_client.meta.result_status == ResultStatus.SKIPPED:
+                # this case applies to upload only
+                return
+            else:
+                try:
+                    logger.debug(f"Finished preparing file {done_client.meta.name}")
+                    tasks = []
+                    for _chunk_id in range(done_client.num_of_chunks):
+                        task = (
+                            asyncio.create_task(done_client.upload_chunk(_chunk_id))
+                            if is_upload
+                            else asyncio.create_task(
+                                done_client.download_chunk(_chunk_id)
+                            )
+                        )
+                        task.add_done_callback(
+                            lambda t, dc=done_client, _chunk_id=_chunk_id: transfer_done_cb(
+                                t, dc, _chunk_id
+                            )
+                        )
+                        tasks.append(task)
+                    await asyncio.gather(*tasks)
+                    await asyncio.gather(*finish_download_upload_tasks)
+                except Exception as error:
+                    done_client.meta.error_details = error
+                    if self._reraise_error_in_file_transfer_work_function:
+                        # Propagate task exceptions to the caller to fail the transfer early.
+                        raise
+
+        def transfer_done_cb(
+            task: asyncio.Task,
+            done_client: SnowflakeStorageClient,
+            chunk_id: int,
+        ) -> None:
+            # Note: chunk_id is 0 based while num_of_chunks is count
+            logger.debug(
+                f"Chunk(id: {chunk_id}) {chunk_id+1}/{done_client.num_of_chunks} of file {done_client.meta.name} reached callback"
+            )
+            if task.exception():
+                done_client.failed_transfers += 1
+                logger.debug(
+                    f"Chunk {chunk_id} of file {done_client.meta.name} failed to transfer for unexpected exception {task.exception()}"
+                )
+            else:
+                done_client.successful_transfers += 1
+            logger.debug(
+                f"Chunk progress: {done_client.meta.name}: completed: {done_client.successful_transfers} failed: {done_client.failed_transfers} total: {done_client.num_of_chunks}"
+            )
+            if (
+                done_client.successful_transfers + done_client.failed_transfers
+                == done_client.num_of_chunks
+            ):
+                if is_upload:
+                    finish_upload_task = asyncio.create_task(
+                        done_client.finish_upload()
+                    )
+                    finish_download_upload_tasks.append(finish_upload_task)
+                    done_client.delete_client_data()
+                else:
+                    finish_download_task = asyncio.create_task(
+                        done_client.finish_download()
+                    )
+                    finish_download_task.add_done_callback(
+                        lambda t, dc=done_client: postprocess_done_cb(t, dc)
+                    )
+                    finish_download_upload_tasks.append(finish_download_task)
+
+        def postprocess_done_cb(
+            task: asyncio.Task,
+            done_client: SnowflakeStorageClient,
+        ) -> None:
+            logger.debug(f"File {done_client.meta.name} reached postprocess callback")
+
+            if task.exception():
+                done_client.failed_transfers += 1
+                logger.debug(
+                    f"File {done_client.meta.name} failed to transfer for unexpected exception {task.exception()}"
+                )
+            # Whether there was an exception or not, we're done the file.
+
+        task_of_files = []
+        for file_client in files:
+            try:
+                # TODO: SNOW-1708819 for code refactoring
+                res = (
+                    await file_client.prepare_upload()
+                    if is_upload
+                    else await file_client.prepare_download()
+                )
+                is_successful = True
+            except Exception as e:
+                res = e
+                file_client.meta.error_details = e
+                is_successful = False
+
+            task = asyncio.create_task(
+                preprocess_done_cb(is_successful, res, done_client=file_client)
+            )
+            task_of_files.append(task)
+        await asyncio.gather(*task_of_files)
+
+        self._results = metas
+
+    async def _transfer_accelerate_config(self) -> None:
+        if self._stage_location_type == S3_FS and self._file_metadata:
+            client = await self._create_file_transfer_client(self._file_metadata[0])
+            self._use_accelerate_endpoint = await client.transfer_accelerate_config()
+
+    async def _create_file_transfer_client(
+        self, meta: SnowflakeFileMeta
+    ) -> SnowflakeStorageClient:
+        if self._stage_location_type == LOCAL_FS:
+            return SnowflakeLocalStorageClient(
+                meta,
+                self._stage_info,
+                4 * megabyte,
+                unsafe_file_write=self._unsafe_file_write,
+            )
+        elif self._stage_location_type == AZURE_FS:
+            return SnowflakeAzureRestClient(
+                meta,
+                self._credentials,
+                AZURE_CHUNK_SIZE,
+                self._stage_info,
+                unsafe_file_write=self._unsafe_file_write,
+            )
+        elif self._stage_location_type == S3_FS:
+            client = SnowflakeS3RestClient(
+                meta=meta,
+                credentials=self._credentials,
+                stage_info=self._stage_info,
+                chunk_size=_chunk_size_calculator(meta.src_file_size),
+                use_accelerate_endpoint=self._use_accelerate_endpoint,
+                use_s3_regional_url=self._use_s3_regional_url,
+                unsafe_file_write=self._unsafe_file_write,
+            )
+            await client.transfer_accelerate_config(self._use_accelerate_endpoint)
+            return client
+        elif self._stage_location_type == GCS_FS:
+            client = SnowflakeGCSRestClient(
+                meta,
+                self._credentials,
+                self._stage_info,
+                self._cursor._connection,
+                self._command,
+                unsafe_file_write=self._unsafe_file_write,
+            )
+            if client.security_token:
+                logger.debug(f"len(GCS_ACCESS_TOKEN): {len(client.security_token)}")
+            else:
+                logger.debug(
+                    "No access token received from GS, requesting presigned url"
+                )
+                await client._update_presigned_url()
+            return client
+        raise Exception(f"{self._stage_location_type} is an unknown stage type")