snowflake-connector-python 3.16.0__cp313-cp313-win_amd64.whl

snowflake/connector/azure_storage_client.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+import base64
+import json
+import os
+import xml.etree.ElementTree as ET
+from datetime import datetime, timezone
+from logging import getLogger
+from random import choice
+from string import hexdigits
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+from .compat import quote
+from .constants import FileHeader, ResultStatus
+from .encryption_util import EncryptionMetadata
+from .storage_client import SnowflakeStorageClient
+from .util_text import get_md5
+from .vendored import requests
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .file_transfer_agent import SnowflakeFileMeta, StorageCredential
+
+logger = getLogger(__name__)
+
+
+class AzureLocation(NamedTuple):
+    container_name: str
+    path: str
+
+
+TOKEN_EXPIRATION_ERR_MESSAGE = (
+    "Signature not valid in the specified time frame",
+    "Server failed to authenticate the request.",
+)
+SFCDIGEST = "x-ms-meta-sfcdigest"
+ENCRYPTION_DATA = "x-ms-meta-encryptiondata"
+MATDESC = "x-ms-meta-matdesc"
+
+
+class SnowflakeAzureRestClient(SnowflakeStorageClient):
+    def __init__(
+        self,
+        meta: SnowflakeFileMeta,
+        credentials: StorageCredential | None,
+        chunk_size: int,
+        stage_info: dict[str, Any],
+        unsafe_file_write: bool = False,
+    ) -> None:
+        super().__init__(
+            meta,
+            stage_info,
+            chunk_size,
+            credentials=credentials,
+            unsafe_file_write=unsafe_file_write,
+        )
+        end_point: str = stage_info["endPoint"]
+        if end_point.startswith("blob."):
+            end_point = end_point[len("blob.") :]
+        self.endpoint = end_point
+        self.storage_account: str = stage_info["storageAccount"]
+        self.azure_location = self.extract_container_name_and_path(
+            stage_info["location"]
+        )
+        self.block_ids: list[str] = []
+
+    @staticmethod
+    def extract_container_name_and_path(stage_location: str) -> AzureLocation:
+        stage_location = os.path.expanduser(stage_location)
+        container_name = stage_location
+        path = ""
+
+        # split stage location as bucket name and path
+        if "/" in stage_location:
+            container_name, _, path = stage_location.partition("/")
+            if path and not path.endswith("/"):
+                path += "/"
+
+        return AzureLocation(container_name=container_name, path=path)
+
+    def _has_expired_token(self, response: requests.Response) -> bool:
+        return response.status_code == 403 and any(
+            message in response.reason for message in TOKEN_EXPIRATION_ERR_MESSAGE
+        )
+
+    def _send_request_with_authentication_and_retry(
+        self,
+        verb: str,
+        url: str,
+        retry_id: int | str,
+        headers: dict[str, Any] = None,
+        data: bytes = None,
+    ) -> requests.Response:
+        if not headers:
+            headers = {}
+
+        def generate_authenticated_url_and_rest_args() -> tuple[bytes, dict[str, Any]]:
+            curtime = datetime.now(timezone.utc).replace(tzinfo=None)
+            timestamp = curtime.strftime("YYYY-MM-DD")
+            sas_token = self.credentials.creds["AZURE_SAS_TOKEN"]
+            if sas_token and sas_token.startswith("?"):
+                sas_token = sas_token[1:]
+            if "?" in url:
+                _url = url + "&" + sas_token
+            else:
+                _url = url + "?" + sas_token
+            headers["Date"] = timestamp
+            rest_args = {"headers": headers}
+            if data:
+                rest_args["data"] = data
+            return _url, rest_args
+
+        return self._send_request_with_retry(
+            verb, generate_authenticated_url_and_rest_args, retry_id
+        )
+
+    def get_file_header(self, filename: str) -> FileHeader | None:
+        """Gets Azure file properties."""
+        container_name = quote(self.azure_location.container_name)
+        path = quote(self.azure_location.path) + quote(filename)
+        meta = self.meta
+        # HTTP HEAD request
+        url = f"https://{self.storage_account}.blob.{self.endpoint}/{container_name}/{path}"
+        retry_id = "HEAD"
+        self.retry_count[retry_id] = 0
+        r = self._send_request_with_authentication_and_retry("HEAD", url, retry_id)
+        if r.status_code == 200:
+            meta.result_status = ResultStatus.UPLOADED
+            enc_data_str = r.headers.get(ENCRYPTION_DATA)
+            encryption_data = None if enc_data_str is None else json.loads(enc_data_str)
+            encryption_metadata = (
+                None
+                if not encryption_data
+                else EncryptionMetadata(
+                    key=encryption_data["WrappedContentKey"]["EncryptedKey"],
+                    iv=encryption_data["ContentEncryptionIV"],
+                    matdesc=r.headers.get(MATDESC),
+                )
+            )
+            return FileHeader(
+                digest=r.headers.get(SFCDIGEST),
+                content_length=int(r.headers.get("Content-Length")),
+                encryption_metadata=encryption_metadata,
+            )
+        elif r.status_code == 404:
+            meta.result_status = ResultStatus.NOT_FOUND_FILE
+            return FileHeader(
+                digest=None, content_length=None, encryption_metadata=None
+            )
+        else:
+            r.raise_for_status()
+
+    def _prepare_file_metadata(self) -> dict[str, str | None]:
+        azure_metadata = {
+            SFCDIGEST: self.meta.sha256_digest,
+        }
+        encryption_metadata = self.encryption_metadata
+        if encryption_metadata:
+            azure_metadata.update(
+                {
+                    ENCRYPTION_DATA: json.dumps(
+                        {
+                            "EncryptionMode": "FullBlob",
+                            "WrappedContentKey": {
+                                "KeyId": "symmKey1",
+                                "EncryptedKey": encryption_metadata.key,
+                                "Algorithm": "AES_CBC_256",
+                            },
+                            "EncryptionAgent": {
+                                "Protocol": "1.0",
+                                "EncryptionAlgorithm": "AES_CBC_128",
+                            },
+                            "ContentEncryptionIV": encryption_metadata.iv,
+                            "KeyWrappingMetadata": {"EncryptionLibrary": "Java 5.3.0"},
+                        }
+                    ),
+                    MATDESC: encryption_metadata.matdesc,
+                }
+            )
+        return azure_metadata
+
+    def _initiate_multipart_upload(self) -> None:
+        self.block_ids = [
+            "".join(choice(hexdigits) for _ in range(20))
+            for _ in range(self.num_of_chunks)
+        ]
+
+    def _upload_chunk(self, chunk_id: int, chunk: bytes) -> None:
+        container_name = quote(self.azure_location.container_name)
+        path = quote(self.azure_location.path + self.meta.dst_file_name.lstrip("/"))
+
+        if self.num_of_chunks > 1:
+            block_id = self.block_ids[chunk_id]
+            url = (
+                f"https://{self.storage_account}.blob.{self.endpoint}/{container_name}/{path}?comp=block"
+                f"&blockid={block_id}"
+            )
+            headers = {"Content-Length": str(len(chunk))}
+            r = self._send_request_with_authentication_and_retry(
+                "PUT", url, chunk_id, headers=headers, data=chunk
+            )
+        else:
+            # single request
+            azure_metadata = self._prepare_file_metadata()
+            url = f"https://{self.storage_account}.blob.{self.endpoint}/{container_name}/{path}"
+            headers = {
+                "x-ms-blob-type": "BlockBlob",
+                "Content-Encoding": "utf-8",
+            }
+            headers.update(azure_metadata)
+            r = self._send_request_with_authentication_and_retry(
+                "PUT", url, chunk_id, headers=headers, data=chunk
+            )
+        r.raise_for_status()  # expect status code 201
+
+    def _complete_multipart_upload(self) -> None:
+        container_name = quote(self.azure_location.container_name)
+        path = quote(self.azure_location.path + self.meta.dst_file_name.lstrip("/"))
+        url = (
+            f"https://{self.storage_account}.blob.{self.endpoint}/{container_name}/{path}?comp"
+            f"=blocklist"
+        )
+        root = ET.Element("BlockList")
+        for block_id in self.block_ids:
+            part = ET.Element("Latest")
+            part.text = block_id
+            root.append(part)
+        # SNOW-1778088: We need to calculate the MD5 sum of this file for Azure Blob storage
+        new_stream = not bool(self.meta.src_stream or self.meta.intermediate_stream)
+        fd = (
+            self.meta.src_stream
+            or self.meta.intermediate_stream
+            or open(self.meta.real_src_file_name, "rb")
+        )
+        try:
+            if not new_stream:
+                # Reset position in file
+                fd.seek(0)
+            file_content = fd.read()
+        finally:
+            if new_stream:
+                fd.close()
+        headers = {
+            "x-ms-blob-content-encoding": "utf-8",
+            "x-ms-blob-content-md5": base64.b64encode(get_md5(file_content)).decode(
+                "utf-8"
+            ),
+        }
+        azure_metadata = self._prepare_file_metadata()
+        headers.update(azure_metadata)
+        retry_id = "COMPLETE"
+        self.retry_count[retry_id] = 0
+        r = self._send_request_with_authentication_and_retry(
+            "PUT", url, "COMPLETE", headers=headers, data=ET.tostring(root)
+        )
+        r.raise_for_status()  # expects status code 201
+
+    def download_chunk(self, chunk_id: int) -> None:
+        container_name = quote(self.azure_location.container_name)
+        path = quote(self.azure_location.path + self.meta.src_file_name.lstrip("/"))
+        url = f"https://{self.storage_account}.blob.{self.endpoint}/{container_name}/{path}"
+        if self.num_of_chunks > 1:
+            chunk_size = self.chunk_size
+            if chunk_id < self.num_of_chunks - 1:
+                _range = f"{chunk_id * chunk_size}-{(chunk_id+1)*chunk_size-1}"
+            else:
+                _range = f"{chunk_id * chunk_size}-"
+            headers = {"Range": f"bytes={_range}"}
+            r = self._send_request_with_authentication_and_retry(
+                "GET", url, chunk_id, headers=headers
+            )  # expect 206
+        else:
+            # single request
+            r = self._send_request_with_authentication_and_retry("GET", url, chunk_id)
+        if r.status_code in (200, 206):
+            self.write_downloaded_chunk(chunk_id, r.content)
+        r.raise_for_status()

snowflake/connector/backoff_policies.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import random
+from typing import Callable, Iterator
+
+"""This module provides common implementations of backoff policies
+
+All backoff policies must be implemented as generator functions with the behaviour specified below. These generator
+functions will be called to create iterators yielding backoff durations.
+
+Args:
+    None
+
+Yields:
+    int: Next backoff duration in seconds
+
+Example:
+    This is an example of a valid backoff policy that always yields a backoff duration of 42 seconds.
+
+    def constant_backoff() -> int:
+        while True
+            yield 42
+
+
+Note:
+    The functions provided in this module are not backoff policies. They are functions returning backoff policies.
+    This is to enable customization of the constants used in backoff computations.
+"""
+
+DEFAULT_BACKOFF_FACTOR = 2
+DEFAULT_BACKOFF_BASE = 1
+DEFAULT_BACKOFF_CAP = 16
+DEFAULT_ENABLE_JITTER = True
+
+
+def mixed_backoff(
+    factor: int = DEFAULT_BACKOFF_FACTOR,
+    base: int = DEFAULT_BACKOFF_BASE,
+    cap: int = DEFAULT_BACKOFF_CAP,
+    enable_jitter: bool = DEFAULT_ENABLE_JITTER,
+) -> Callable[..., Iterator[int]]:
+    """Randomly chooses between exponential and constant backoff. Uses equal jitter.
+
+    Args:
+        factor (int): Exponential base for the exponential term.
+        base (int): Initial backoff time in seconds. Constant coefficient for the exponential term.
+        cap (int): Maximum backoff time in seconds.
+        enable_jitter (int): Whether to enable equal jitter on computed durations. For details see
+            https://www.awsarchitectureblog.com/2015/03/backoff.html
+
+    Returns:
+        Callable: generator function implementing the mixed backoff policy
+    """
+
+    def generator():
+        cnt = 0
+        sleep = base
+
+        yield sleep
+        while True:
+            cnt += 1
+
+            # equal jitter
+            mult_factor = random.choice([-1, 1])
+            jitter_amount = 0.5 * sleep * mult_factor if enable_jitter else 0
+            sleep = int(
+                random.choice(
+                    [sleep + jitter_amount, base * factor**cnt + jitter_amount]
+                )
+            )
+            sleep = min(cap, sleep)
+
+            yield sleep
+
+    return generator
+
+
+def linear_backoff(
+    factor: int = DEFAULT_BACKOFF_FACTOR,
+    base: int = DEFAULT_BACKOFF_BASE,
+    cap: int = DEFAULT_BACKOFF_CAP,
+    enable_jitter: bool = DEFAULT_ENABLE_JITTER,
+) -> Callable[..., Iterator[int]]:
+    """Standard linear backoff. Uses full jitter.
+
+    Args:
+        factor (int): Linear increment every iteration.
+        base (int): Initial backoff time in seconds.
+        cap (int): Maximum backoff time in seconds.
+        enable_jitter (int): Whether to enable full jitter on computed durations. For details see
+            https://www.awsarchitectureblog.com/2015/03/backoff.html
+
+    Returns:
+        Callable: generator function implementing the linear backoff policy
+    """
+
+    def generator():
+        sleep = base
+
+        yield sleep
+        while True:
+            sleep += factor
+            sleep = min(cap, sleep)
+
+            # full jitter
+            yield random.randint(0, sleep) if enable_jitter else sleep
+
+    return generator
+
+
+def exponential_backoff(
+    factor: int = DEFAULT_BACKOFF_FACTOR,
+    base: int = DEFAULT_BACKOFF_BASE,
+    cap: int = DEFAULT_BACKOFF_CAP,
+    enable_jitter: bool = DEFAULT_ENABLE_JITTER,
+) -> Callable[..., Iterator[int]]:
+    """Standard exponential backoff. Uses full jitter.
+
+    Args:
+        factor (int): Exponential base for the exponential term.
+        base (int): Initial backoff time in seconds. Constant coefficient for the exponential term.
+        cap (int): Maximum backoff time in seconds.
+        enable_jitter (int): Whether to enable full jitter on computed durations. For details see
+            https://www.awsarchitectureblog.com/2015/03/backoff.html
+
+    Returns:
+        Callable: generator function implementing the exponential backoff policy
+    """
+
+    def generator():
+        sleep = base
+
+        yield sleep
+        while True:
+            sleep *= factor
+            sleep = min(cap, sleep)
+
+            # full jitter
+            yield random.randint(0, sleep) if enable_jitter else sleep
+
+    return generator

snowflake/connector/bind_upload_agent.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import os
+import uuid
+from io import BytesIO
+from logging import getLogger
+from typing import TYPE_CHECKING
+
+from ._utils import (
+    _PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS_STRING,
+    get_temp_type_for_object,
+)
+from .errors import BindUploadError, Error
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .cursor import SnowflakeCursor
+
+logger = getLogger(__name__)
+
+
+class BindUploadAgent:
+
+    def __init__(
+        self,
+        cursor: SnowflakeCursor,
+        rows: list[bytes],
+        stream_buffer_size: int = 1024 * 1024 * 10,
+    ) -> None:
+        """Construct an agent that uploads binding parameters as CSV files to a temporary stage.
+
+        Args:
+            cursor: The cursor object.
+            rows: Rows of binding parameters in CSV format.
+            stream_buffer_size: Size of each file, default to 10MB.
+        """
+        self._use_scoped_temp_object = (
+            cursor.connection._session_parameters.get(
+                _PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS_STRING, False
+            )
+            if cursor.connection._session_parameters
+            else False
+        )
+        self._STAGE_NAME = (
+            "SNOWPARK_TEMP_STAGE_BIND" if self._use_scoped_temp_object else "SYSTEMBIND"
+        )
+        self.cursor = cursor
+        self.rows = rows
+        self._stream_buffer_size = stream_buffer_size
+        self.stage_path = f"@{self._STAGE_NAME}/{uuid.uuid4().hex}"
+
+    def _create_stage(self) -> None:
+        create_stage_sql = (
+            f"create or replace {get_temp_type_for_object(self._use_scoped_temp_object)} stage {self._STAGE_NAME} "
+            "file_format=(type=csv field_optionally_enclosed_by='\"')"
+        )
+        self.cursor.execute(create_stage_sql)
+
+    def upload(self) -> None:
+        try:
+            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:
+                f.seek(0)
+                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.")
+                raise BindUploadError from err
+            f.close()