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

snowflake/connector/dbapi.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+"""This module implements some constructors and singletons as required by the DB API v2.0 (PEP-249)."""
+
+from __future__ import annotations
+
+import datetime
+import time
+
+from .constants import (
+    get_binary_types,
+    get_number_types,
+    get_string_types,
+    get_timestamp_types,
+)
+
+
+class _DBAPITypeObject:
+    def __init__(self, *values) -> None:
+        self.values = values
+
+    def __cmp__(self, other):
+        if other in self.values:
+            return 0
+        if other < self.values:
+            return 1
+        else:
+            return -1
+
+
+Date = datetime.date
+Time = datetime.time
+Timestamp = datetime.datetime
+
+
+def DateFromTicks(ticks: float) -> datetime.date:
+    return Date(*time.localtime(ticks)[:3])
+
+
+def TimeFromTicks(ticks: float) -> datetime.time:
+    return Time(*time.localtime(ticks)[3:6])
+
+
+def TimestampFromTicks(ticks: float) -> datetime.datetime:
+    return Timestamp(*time.localtime(ticks)[:6])
+
+
+Binary = bytes
+
+STRING = _DBAPITypeObject(get_string_types())
+BINARY = _DBAPITypeObject(get_binary_types())
+NUMBER = _DBAPITypeObject(get_number_types())
+DATETIME = _DBAPITypeObject(get_timestamp_types())
+ROWID = _DBAPITypeObject()

snowflake/connector/description.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python
+"""Various constants."""
+
+from __future__ import annotations
+
+import platform
+import sys
+
+from .version import VERSION
+
+SNOWFLAKE_CONNECTOR_VERSION = ".".join(str(v) for v in VERSION[0:3])
+PYTHON_VERSION = ".".join(str(v) for v in sys.version_info[:3])
+OPERATING_SYSTEM = platform.system()
+PLATFORM = platform.platform()
+IMPLEMENTATION = platform.python_implementation()
+COMPILER = platform.python_compiler()
+
+CLIENT_NAME = "PythonConnector"  # don't change!
+CLIENT_VERSION = ".".join([str(v) for v in VERSION[:3]])

snowflake/connector/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
+    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: SnowflakeConnection):
+        self._connection = connection
+
+    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}")
+
+        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 cursor._execute_helper(
+                sql, binding_params=processed_params, is_internal=True
+            )
+
+
+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")

snowflake/connector/encryption_util.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import base64
+import json
+import os
+import tempfile
+from logging import getLogger
+from typing import IO, TYPE_CHECKING
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+from .compat import PKCS5_OFFSET, PKCS5_PAD, PKCS5_UNPAD
+from .constants import UTF8, EncryptionMetadata, MaterialDescriptor, kilobyte
+from .file_util import owner_rw_opener
+from .util_text import random_string
+
+block_size = int(algorithms.AES.block_size / 8)  # in bytes
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .storage_client import SnowflakeFileEncryptionMaterial
+
+logger = getLogger(__name__)
+
+
+def matdesc_to_unicode(matdesc: MaterialDescriptor) -> str:
+    """Convert Material Descriptor to Unicode String."""
+    return str(
+        json.dumps(
+            {
+                "queryId": matdesc.query_id,
+                "smkId": str(matdesc.smk_id),
+                "keySize": str(matdesc.key_size),
+            },
+            separators=(",", ":"),
+        )
+    )
+
+
+class SnowflakeEncryptionUtil:
+    @staticmethod
+    def get_secure_random(byte_length: int) -> bytes:
+        return os.urandom(byte_length)
+
+    @staticmethod
+    def encrypt_stream(
+        encryption_material: SnowflakeFileEncryptionMaterial,
+        src: IO[bytes],
+        out: IO[bytes],
+        chunk_size: int = 64 * kilobyte,  # block_size * 4 * 1024,
+    ) -> EncryptionMetadata:
+        """Reads content from src and write the encrypted content into out.
+
+        This function is sensitive to current position of src and out.
+        It does not seek to position 0 in neither stream objects before or after the encryption.
+
+        Args:
+            encryption_material: The encryption material for file.
+            src: The input stream.
+            out: The output stream.
+            chunk_size: The size of read chunks (Default value = block_size * 4 * 1024
+
+        Returns:
+            The encryption metadata.
+        """
+        logger = getLogger(__name__)
+        decoded_key = base64.standard_b64decode(
+            encryption_material.query_stage_master_key
+        )
+        key_size = len(decoded_key)
+        logger.debug("key_size = %s", key_size)
+
+        # Generate key for data encryption
+        iv_data = SnowflakeEncryptionUtil.get_secure_random(block_size)
+        file_key = SnowflakeEncryptionUtil.get_secure_random(key_size)
+        backend = default_backend()
+        cipher = Cipher(algorithms.AES(file_key), modes.CBC(iv_data), backend=backend)
+        encryptor = cipher.encryptor()
+
+        padded = False
+        while True:
+            chunk = src.read(chunk_size)
+            if len(chunk) == 0:
+                break
+            elif len(chunk) % block_size != 0:
+                chunk = PKCS5_PAD(chunk, block_size)
+                padded = True
+            out.write(encryptor.update(chunk))
+        if not padded:
+            out.write(encryptor.update(block_size * chr(block_size).encode(UTF8)))
+        out.write(encryptor.finalize())
+
+        # encrypt key with QRMK
+        cipher = Cipher(algorithms.AES(decoded_key), modes.ECB(), backend=backend)
+        encryptor = cipher.encryptor()
+        enc_kek = (
+            encryptor.update(PKCS5_PAD(file_key, block_size)) + encryptor.finalize()
+        )
+
+        mat_desc = MaterialDescriptor(
+            smk_id=encryption_material.smk_id,
+            query_id=encryption_material.query_id,
+            key_size=key_size * 8,
+        )
+        metadata = EncryptionMetadata(
+            key=base64.b64encode(enc_kek).decode("utf-8"),
+            iv=base64.b64encode(iv_data).decode("utf-8"),
+            matdesc=matdesc_to_unicode(mat_desc),
+        )
+        return metadata
+
+    @staticmethod
+    def encrypt_file(
+        encryption_material: SnowflakeFileEncryptionMaterial,
+        in_filename: str,
+        chunk_size: int = 64 * kilobyte,
+        tmp_dir: str | None = None,
+    ) -> tuple[EncryptionMetadata, str]:
+        """Encrypts a file in a temporary directory.
+
+        Args:
+            encryption_material: The encryption material for file.
+            in_filename: The input file's name.
+            chunk_size: The size of read chunks (Default value = block_size * 4 * 1024).
+            tmp_dir: Temporary directory to use, optional (Default value = None).
+
+        Returns:
+            The encryption metadata and the encrypted file's location.
+        """
+        logger = getLogger(__name__)
+        temp_output_fd, temp_output_file = tempfile.mkstemp(
+            text=False, dir=tmp_dir, prefix=os.path.basename(in_filename) + "#"
+        )
+        logger.debug(
+            "unencrypted file: %s, temp file: %s, tmp_dir: %s",
+            in_filename,
+            temp_output_file,
+            tmp_dir,
+        )
+        with open(in_filename, "rb") as infile:
+            with os.fdopen(temp_output_fd, "wb") as outfile:
+                metadata = SnowflakeEncryptionUtil.encrypt_stream(
+                    encryption_material, infile, outfile, chunk_size
+                )
+        return metadata, temp_output_file
+
+    @staticmethod
+    def decrypt_stream(
+        metadata: EncryptionMetadata,
+        encryption_material: SnowflakeFileEncryptionMaterial,
+        src: IO[bytes],
+        out: IO[bytes],
+        chunk_size: int = 64 * kilobyte,  # block_size * 4 * 1024,
+    ) -> None:
+        """To read from `src` stream then decrypt to `out` stream."""
+
+        key_base64 = metadata.key
+        iv_base64 = metadata.iv
+        decoded_key = base64.standard_b64decode(
+            encryption_material.query_stage_master_key
+        )
+        key_bytes = base64.standard_b64decode(key_base64)
+        iv_bytes = base64.standard_b64decode(iv_base64)
+
+        backend = default_backend()
+        cipher = Cipher(algorithms.AES(decoded_key), modes.ECB(), backend=backend)
+        decryptor = cipher.decryptor()
+        file_key = PKCS5_UNPAD(decryptor.update(key_bytes) + decryptor.finalize())
+        cipher = Cipher(algorithms.AES(file_key), modes.CBC(iv_bytes), backend=backend)
+        decryptor = cipher.decryptor()
+
+        last_decrypted_chunk = None
+        chunk = src.read(chunk_size)
+        while len(chunk) != 0:
+            if last_decrypted_chunk is not None:
+                out.write(last_decrypted_chunk)
+            d = decryptor.update(chunk)
+            last_decrypted_chunk = d
+            chunk = src.read(chunk_size)
+
+        if last_decrypted_chunk is not None:
+            offset = PKCS5_OFFSET(last_decrypted_chunk)
+            out.write(last_decrypted_chunk[:-offset])
+        out.write(decryptor.finalize())
+
+    @staticmethod
+    def decrypt_file(
+        metadata: EncryptionMetadata,
+        encryption_material: SnowflakeFileEncryptionMaterial,
+        in_filename: str,
+        chunk_size: int = 64 * kilobyte,
+        tmp_dir: str | None = None,
+        unsafe_file_write: bool = False,
+    ) -> str:
+        """Decrypts a file and stores the output in the temporary directory.
+
+        Args:
+            metadata: The file's metadata input.
+            encryption_material: The file's encryption material.
+            in_filename: The name of the input file.
+            chunk_size: The size of read chunks (Default value = block_size * 4 * 1024).
+            tmp_dir: Temporary directory to use, optional (Default value = None).
+
+        Returns:
+            The decrypted file's location.
+        """
+        temp_output_file = f"{os.path.basename(in_filename)}#{random_string()}"
+        if tmp_dir:
+            temp_output_file = os.path.join(tmp_dir, temp_output_file)
+
+        logger.debug("encrypted file: %s, tmp file: %s", in_filename, temp_output_file)
+
+        file_opener = None if unsafe_file_write else owner_rw_opener
+        with open(in_filename, "rb") as infile:
+            with open(temp_output_file, "wb", opener=file_opener) as outfile:
+                SnowflakeEncryptionUtil.decrypt_stream(
+                    metadata, encryption_material, infile, outfile, chunk_size
+                )
+        return temp_output_file

snowflake/connector/errorcode.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+# network
+ER_FAILED_TO_CONNECT_TO_DB = 250001
+ER_CONNECTION_IS_CLOSED = 250002
+ER_FAILED_TO_REQUEST = 250003
+ER_NOT_HTTPS_USED = 250004
+ER_FAILED_TO_SERVER = 250005
+ER_IDP_CONNECTION_ERROR = 250006
+ER_INCORRECT_DESTINATION = 250007
+ER_UNABLE_TO_OPEN_BROWSER = 250008
+ER_UNABLE_TO_START_WEBSERVER = 250009
+ER_INVALID_CERTIFICATE = 250011  # not used but keep here to reserve errno
+ER_INVALID_BACKOFF_POLICY = 250012
+
+# connection
+ER_NO_ACCOUNT_NAME = 251001
+ER_OLD_PYTHON = 251002
+ER_NO_WINDOWS_SUPPORT = 251003
+ER_FAILED_TO_GET_BOOTSTRAP = 251004
+ER_NO_USER = 251005
+ER_NO_PASSWORD = 251006
+ER_INVALID_VALUE = 251007
+ER_INVALID_PRIVATE_KEY = 251008
+ER_NO_HOSTNAME_FOUND = 251009
+ER_JWT_RETRY_EXPIRED = 251010
+ER_CONNECTION_TIMEOUT = 251011
+ER_RETRYABLE_CODE = 251012
+ER_NO_CLIENT_ID = 251013
+ER_OAUTH_STATE_CHANGED = 251014
+ER_OAUTH_CALLBACK_ERROR = 251015
+ER_OAUTH_SERVER_TIMEOUT = 251016
+ER_INVALID_WIF_SETTINGS = 251017
+ER_WIF_CREDENTIALS_NOT_FOUND = 251018
+ER_EXPERIMENTAL_AUTHENTICATION_NOT_SUPPORTED = 251019
+ER_NO_CLIENT_SECRET = 251020
+
+# cursor
+ER_FAILED_TO_REWRITE_MULTI_ROW_INSERT = 252001
+ER_NO_ADDITIONAL_CHUNK = 252002
+ER_NOT_POSITIVE_SIZE = 252003
+ER_FAILED_PROCESSING_PYFORMAT = 252004
+ER_FAILED_TO_CONVERT_ROW_TO_PYTHON_TYPE = 252005
+ER_CURSOR_IS_CLOSED = 252006
+ER_FAILED_TO_RENEW_SESSION = 252007
+ER_UNSUPPORTED_METHOD = 252008
+ER_NO_DATA_FOUND = 252009
+ER_CHUNK_DOWNLOAD_FAILED = 252010
+ER_NOT_IMPLICITY_SNOWFLAKE_DATATYPE = 252011
+ER_FAILED_PROCESSING_QMARK = 252012
+
+# file_transfer
+ER_INVALID_STAGE_FS = 253001
+ER_FAILED_TO_DOWNLOAD_FROM_STAGE = 253002
+ER_FAILED_TO_UPLOAD_TO_STAGE = 253003
+ER_INVALID_STAGE_LOCATION = 253004
+ER_LOCAL_PATH_NOT_DIRECTORY = 253005
+ER_FILE_NOT_EXISTS = 253006
+ER_COMPRESSION_NOT_SUPPORTED = 253007
+ER_INTERNAL_NOT_MATCH_ENCRYPT_MATERIAL = 253008
+ER_FAILED_TO_CHECK_EXISTING_FILES = 253009
+
+# ocsp
+ER_OCSP_URL_INFO_MISSING = 254001
+ER_OCSP_RESPONSE_UNAVAILABLE = 254002
+ER_OCSP_RESPONSE_FETCH_EXCEPTION = 254003
+ER_OCSP_FAILED_TO_CONNECT_CACHE_SERVER = 254004
+ER_OCSP_RESPONSE_CERT_STATUS_INVALID = 254005
+ER_OCSP_RESPONSE_CERT_STATUS_UNKNOWN = 254006
+ER_OCSP_RESPONSE_CERT_STATUS_REVOKED = 254007
+ER_OCSP_RESPONSE_STATUS_UNSUCCESSFUL = 254008
+ER_OCSP_RESPONSE_ATTACHED_CERT_INVALID = 254009
+ER_OCSP_RESPONSE_ATTACHED_CERT_EXPIRED = 254010
+ER_OCSP_RESPONSE_INVALID_SIGNATURE = 254011
+ER_OCSP_RESPONSE_INVALID_EXPIRY_INFO_MISSING = 254012
+ER_OCSP_RESPONSE_EXPIRED = 254013
+ER_OCSP_RESPONSE_FETCH_FAILURE = 254014
+ER_OCSP_RESPONSE_LOAD_FAILURE = 254015
+ER_OCSP_RESPONSE_CACHE_DOWNLOAD_FAILED = 254016
+ER_OCSP_RESPONSE_CACHE_DECODE_FAILED = 254017
+ER_INVALID_OCSP_RESPONSE_SSD = 254018
+ER_INVALID_SSD = 254019
+
+# converter
+ER_NOT_SUPPORT_DATA_TYPE = 255001
+ER_NO_PYARROW = 255002
+ER_NO_ARROW_RESULT = 255003
+ER_NO_PYARROW_SNOWSQL = 255004
+ER_FAILED_TO_READ_ARROW_STREAM = 255005
+ER_NO_NUMPY = 255006