snowflake-connector-python 5.2.0__cp315-cp315-win_amd64.whl

snowflake/connector/_internal/protobuf_gen/proto_exception.py

@@ -0,0 +1,17 @@
+class ProtoException(Exception):
+    pass
+
+class ProtoApplicationException(ProtoException):
+    def __init__(self, error, *args, **kwargs):
+        self.api_error_pb = error
+        super().__init__(error, *args, **kwargs)
+
+    def __getattr__(self, name):
+        try:
+            return getattr(self.api_error_pb, name)
+        except AttributeError:
+            raise AttributeError(f"'{type(self).__name__}' underlying protobuf object has no attribute '{name}'")
+
+
+class ProtoTransportException(ProtoException):
+    pass

snowflake/connector/_internal/result_batch.py

@@ -0,0 +1,150 @@
+"""Shared result-batch mixin for sync and async result batch implementations."""
+
+from __future__ import annotations
+
+from enum import Enum, unique
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from .arrow_context import ArrowConverterContext
+from .errorhandler import ErrorHandlerMixin
+from .protobuf_gen.database_driver_v1_pb2 import ColumnMetadata, ResultChunk
+
+
+if TYPE_CHECKING:
+    from ..aio.connection import Connection as AsyncConnection
+    from ..connection import Connection
+    from .cursor.result_metadata import ResultMetadata
+
+
+_ResultBatchT = TypeVar("_ResultBatchT", bound="ResultBatchMixin")
+
+
+@unique
+class IterUnit(Enum):
+    """Controls what ``ResultBatch.create_iter`` yields."""
+
+    ROW_UNIT = "row"
+    TABLE_UNIT = "table"
+
+    @classmethod
+    def of(cls, value: IterUnit | str) -> IterUnit:
+        if isinstance(value, cls):
+            return value
+        return cls(value)
+
+
+@unique
+class IterTableStructure(Enum):
+    """Controls what table format ``TABLE_UNIT`` iteration produces."""
+
+    ARROW = "arrow"
+    PANDAS = "pandas"
+
+    @classmethod
+    def of(cls, value: IterTableStructure | str) -> IterTableStructure:
+        if isinstance(value, cls):
+            return value
+        return cls(value)
+
+
+class ResultBatchMixin(ErrorHandlerMixin):
+    """Zero-I/O result-batch members shared by sync and async batch classes."""
+
+    _chunk: ResultChunk
+    _description: list[ResultMetadata]
+    _connection: Connection | AsyncConnection | None
+    _columns: list[ColumnMetadata]
+    _arrow_context: ArrowConverterContext
+    _numpy: bool
+    _use_core_arrow: bool
+    _arrow_stream_ptr: int | None
+
+    def __init__(
+        self,
+        chunk: ResultChunk,
+        description: list[ResultMetadata],
+        connection: Connection | AsyncConnection | None,
+        columns: list[ColumnMetadata] | None = None,
+    ) -> None:
+        self._chunk = chunk
+        self._description = description
+        self._connection = connection
+        self._columns = list(columns) if columns else []
+        self._arrow_context = (
+            ArrowConverterContext.create(connection) if connection is not None else ArrowConverterContext()
+        )
+        self._numpy = bool(connection.config.numpy) if connection is not None else False
+        self._use_core_arrow = bool(connection.config.use_core_arrow) if connection is not None else False
+        self._arrow_stream_ptr = None
+
+    @classmethod
+    def from_chunks(
+        cls: type[_ResultBatchT],
+        chunks: list[ResultChunk] | None,
+        description: list[ResultMetadata] | None,
+        connection: Connection | AsyncConnection | None,
+        columns: list[ColumnMetadata] | None = None,
+    ) -> list[_ResultBatchT] | None:
+        """Create a list of batches from raw result chunks, or ``None`` if unavailable."""
+        if chunks is None or description is None:
+            return None
+        return [cls(chunk=chunk, description=description, connection=connection, columns=columns) for chunk in chunks]
+
+    # ------------------------------------------------------------------
+    # Properties
+    # ------------------------------------------------------------------
+
+    @property
+    def rowcount(self) -> int:
+        return self._chunk.row_count
+
+    @property
+    def compressed_size(self) -> int | None:
+        if self._chunk.HasField("remote"):
+            return self._chunk.remote.compressed_size
+        return None
+
+    @property
+    def uncompressed_size(self) -> int | None:
+        if self._chunk.HasField("remote"):
+            return self._chunk.remote.uncompressed_size
+        return None
+
+    @property
+    def column_names(self) -> list[str]:
+        return [col.name for col in self._description]
+
+    @property
+    def _errorhandler_connection(self) -> Connection | AsyncConnection | None:
+        return self._connection
+
+    # ------------------------------------------------------------------
+    # Pickle support
+    # ------------------------------------------------------------------
+
+    def __getstate__(self) -> dict[str, Any]:
+        return {
+            "chunk_bytes": self._chunk.SerializeToString(),
+            "description": self._description,
+            "column_bytes": [c.SerializeToString() for c in self._columns],
+            "arrow_context": self._arrow_context,
+            "numpy": self._numpy,
+            "use_core_arrow": self._use_core_arrow,
+        }
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        chunk = ResultChunk()
+        chunk.ParseFromString(state["chunk_bytes"])
+        self._chunk = chunk
+        self._description = state["description"]
+        columns: list[ColumnMetadata] = []
+        for raw in state.get("column_bytes", []):
+            col = ColumnMetadata()
+            col.ParseFromString(raw)
+            columns.append(col)
+        self._columns = columns
+        self._arrow_context = state.get("arrow_context", ArrowConverterContext())
+        self._numpy = state.get("numpy", False)
+        self._use_core_arrow = state.get("use_core_arrow", False)
+        self._connection = None
+        self._arrow_stream_ptr = None

snowflake/connector/_internal/snowflake_restful.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import json
+
+from typing import TYPE_CHECKING
+from urllib.parse import urlparse
+
+from snowflake.connector._internal.api_client.client_api import core_driver
+from snowflake.connector._internal.protobuf_gen.database_driver_v1_pb2 import (
+    TOKEN_REQUEST_TYPE_ISSUE,
+    TOKEN_REQUEST_TYPE_RENEW,
+    ConnectionGetInfoResponse,
+)
+from snowflake.connector.errors import OperationalError
+
+
+if TYPE_CHECKING:
+    from ..connection import SnowflakeConnection
+
+DEFAULT_PROTOCOL = "https"
+DEFAULT_PORT = 443
+APPLICATION_JSON = "application/json"
+
+_TOKEN_REQUEST_TYPE_MAP = {
+    "ISSUE": TOKEN_REQUEST_TYPE_ISSUE,
+    "RENEW": TOKEN_REQUEST_TYPE_RENEW,
+}
+
+
+class SnowflakeRestful:
+    """Backward-compatible REST interface wrapping the driver's internal gRPC API.
+
+    Used by Snowflake CLI to issue Snowflake REST API calls (v1 and v2)
+    through the driver's TLS-configured HTTP client.
+    """
+
+    def __init__(self, connection: SnowflakeConnection) -> None:
+        self._connection = connection
+
+    @property
+    def _connection_info(self) -> ConnectionGetInfoResponse:
+        return self._connection._get_connection_info()
+
+    def get_user_agent(self) -> str | None:
+        """User-Agent string built by the driver core (None before connect).
+
+        The Rust core constructs the User-Agent from the client app id, version
+        and platform; consumers building REST requests should use this rather
+        than reconstructing the value themselves.
+        """
+        user_agent: str | None = self._connection_info.user_agent
+        return user_agent if user_agent else None
+
+    @property
+    def token(self) -> str | None:
+        """Required by Python API. Returns None after close (handle released)."""
+        if self._connection.is_closed():
+            return None
+        session_token: str | None = self._connection_info.session_token
+        return session_token if session_token else None
+
+    @property
+    def _host(self) -> str | None:
+        host: str | None = self._connection_info.host
+        return host
+
+    @property
+    def _protocol(self) -> str | None:
+        url = self._connection_info.server_url
+        if not url:
+            return DEFAULT_PROTOCOL
+        return urlparse(url).scheme or DEFAULT_PROTOCOL
+
+    @property
+    def _port(self) -> int | None:
+        info = self._connection_info
+        if info.port:
+            return info.port
+        url = info.server_url
+        if not url:
+            return DEFAULT_PORT
+        return urlparse(url).port or DEFAULT_PORT
+
+    @property
+    def server_url(self) -> str:
+        url = self._connection_info.server_url
+        if url:
+            url = url.rstrip("/")
+            parsed = urlparse(url)
+            if parsed.port is None:
+                port = self._port
+                return f"{parsed.scheme}://{parsed.hostname}:{port}{parsed.path}"
+            return url
+        host = self._host
+        if not host:
+            raise ValueError("server_url and host are both unavailable from connection info")
+        return f"{self._protocol}://{host}:{self._port}"
+
+    @property
+    def master_token(self) -> str | None:
+        """Returns None after close (handle released)."""
+        if self._connection.is_closed():
+            return None
+        info = self._connection._get_connection_info(include_master_token=True)
+        master_token: str = info.master_token
+        return master_token if master_token else None
+
+    def _send_http(
+        self,
+        method: str,
+        url: str,
+        headers: dict[str, str] | None = None,
+        body: dict | str | bytes | None = None,
+    ) -> tuple[int, dict[str, str], bytes]:
+        if url.startswith(("http://", "https://")):
+            parsed = urlparse(url)
+            server_parsed = urlparse(self.server_url)
+            if parsed.hostname and server_parsed.hostname and parsed.hostname != server_parsed.hostname:
+                raise ValueError(
+                    f"Absolute URL host '{parsed.hostname}' does not match "
+                    f"connection server host '{server_parsed.hostname}'. "
+                    f"Pass a relative path instead."
+                )
+            relative_url = parsed.path
+            if parsed.query:
+                relative_url = f"{relative_url}?{parsed.query}"
+            url = relative_url
+
+        if body is not None:
+            if isinstance(body, dict):
+                encoded_body = json.dumps(body).encode("utf-8")
+            elif isinstance(body, str):
+                encoded_body = body.encode("utf-8")
+            else:
+                encoded_body = body
+        else:
+            encoded_body = None
+
+        response = core_driver.connection_send_http(
+            conn_handle=self._connection.conn_handle,  # type: ignore[arg-type]
+            method=method,
+            url=url,
+            headers=headers or {},
+            body=encoded_body,
+        )
+        return response.status_code, dict(response.headers), response.body
+
+    _SUPPORTED_REQUEST_METHODS = {"get", "post"}
+
+    def request(
+        self,
+        url: str,
+        body: dict | None = None,
+        method: str = "post",
+        client: str = "sfsql",
+        timeout: int | None = None,
+        _no_results: bool = False,
+        _include_retry_params: bool = False,
+        _no_retry: bool = False,
+    ) -> dict:
+        if client == "sfsql":
+            raise NotImplementedError(
+                "request() with client='sfsql' is not supported. Use cursor.execute() for SQL queries."
+            )
+
+        if method.lower() not in self._SUPPORTED_REQUEST_METHODS:
+            raise ValueError(f"Unsupported HTTP method '{method}'. Only 'get' and 'post' are supported.")
+
+        headers = {
+            "Content-Type": APPLICATION_JSON,
+            "Accept": APPLICATION_JSON,
+        }
+
+        status_code, _, response_body = self._send_http(
+            method=method,
+            url=url,
+            headers=headers,
+            body=body,
+        )
+
+        if status_code >= 400:
+            raise OperationalError.from_http_response(status_code, response_body)
+
+        return json.loads(response_body) if response_body else {}
+
+    def fetch(
+        self,
+        method: str,
+        full_url: str,
+        headers: dict[str, str],
+        data: str | None = None,
+        timeout: int | None = None,
+        token: str | None = None,
+        no_retry: bool = False,
+        raise_raw_http_failure: bool = False,
+        **kwargs: object,
+    ) -> dict | list:
+        status_code, _, response_body = self._send_http(
+            method=method,
+            url=full_url,
+            headers=headers,
+            body=data,
+        )
+
+        if raise_raw_http_failure and status_code >= 400:
+            raise OperationalError.from_http_response(status_code, response_body)
+
+        return json.loads(response_body) if response_body else {}
+
+    def _token_request(self, request_type: str) -> dict:
+        proto_type = _TOKEN_REQUEST_TYPE_MAP.get(request_type.upper())
+        if proto_type is None:
+            raise ValueError(f"Unknown token request type: {request_type!r}. Must be 'ISSUE' or 'RENEW'.")
+
+        response = core_driver.connection_request_token(
+            conn_handle=self._connection.conn_handle,  # type: ignore[arg-type]
+            request_type=proto_type,
+        )
+        data: dict = {"sessionToken": response.session_token}
+        if response.HasField("validity_in_seconds"):
+            data["validityInSecondsST"] = response.validity_in_seconds
+        return {"data": data, "success": True}

snowflake/connector/_internal/snowpark_backcompat_helpers.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+_STAGE_PREFIXES = ("@", "snow://", "/")
+
+
+def file_uri(path: str) -> str:
+    """Add a ``file://`` prefix unless *path* already has one or is quoted."""
+    stripped = path.strip()
+    return stripped if stripped.startswith(("'", "file://")) else f"file://{stripped}"
+
+
+def stage_ref(path: str) -> str:
+    """Add an ``@`` prefix unless *path* already names a stage."""
+    stripped = path.strip()
+    return stripped if stripped.startswith(("'", *_STAGE_PREFIXES)) else f"@{stripped}"
+
+
+def unquoted_stage_ref(path: str) -> str:
+    """Add an ``@`` prefix and strip quoting; the streaming RPC takes the stage name as a field, not SQL."""
+    stripped = path.strip()
+    if len(stripped) > 1 and stripped.startswith("'") and stripped.endswith("'"):
+        stripped = stripped[1:-1]
+    return stripped if stripped.startswith(_STAGE_PREFIXES) else f"@{stripped}"
+
+
+def put_get_options(options: dict[str, Any] | None) -> str:
+    """Render an options mapping as trailing ``KEY=VALUE`` pairs, unquoted."""
+    return " ".join(f"{name}={value}" for name, value in (options or {}).items())

snowflake/connector/_internal/sqlstate.py

@@ -0,0 +1,11 @@
+"""
+SQLSTATE codes.
+
+It should be aligned with the reference snowflake-connector-python.
+"""
+
+from __future__ import annotations
+
+
+SQLSTATE_SUCCESS = "00000"
+SQLSTATE_CONNECTION_NOT_EXISTS = "08003"

snowflake/connector/_internal/statement_utils.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, Generator
+from contextlib import asynccontextmanager, contextmanager
+
+from .api_client.client_api import async_core_driver, core_driver
+from .protobuf_gen.database_driver_v1_pb2 import (
+    ConnectionHandle,
+    DatabaseFetchChunkResponse,
+    PrepareResult,
+    ResultSetDescriptor,
+    ResultSetGetStreamResponse,
+    StatementHandle,
+)
+from .sqlstate import SQLSTATE_SUCCESS
+
+
+@contextmanager
+def statement(conn_handle: ConnectionHandle, query: str) -> Generator[StatementHandle]:
+    """Context manager that owns the full lifecycle of a statement handle.
+
+    Allocates a new statement on the server, binds the given SQL query to it,
+    and yields the ``StatementHandle`` for execution.  The statement is
+    guaranteed to be released when the context exits, even if an exception
+    is raised.
+
+    Args:
+        conn_handle: Active Snowflake conn_handle.
+        query: SQL text to bind to the newly created statement.
+
+    Yields:
+        StatementHandle: A handle that can be passed to ``statement_execute``
+        or other statement-level APIs.
+    """
+    stmt_handle = core_driver.statement_new(conn_handle=conn_handle).stmt_handle
+    try:
+        core_driver.statement_set_query(stmt_handle=stmt_handle, query=query)
+        yield stmt_handle
+    finally:
+        core_driver.statement_release(stmt_handle=stmt_handle)
+
+
+@asynccontextmanager
+async def async_statement(conn_handle: ConnectionHandle, query: str) -> AsyncGenerator[StatementHandle]:
+    stmt_handle = (await async_core_driver.statement_new(conn_handle=conn_handle)).stmt_handle
+    try:
+        await async_core_driver.statement_set_query(stmt_handle=stmt_handle, query=query)
+        yield stmt_handle
+    finally:
+        await async_core_driver.statement_release(stmt_handle=stmt_handle)
+
+
+def get_stream_ptr(result: DatabaseFetchChunkResponse | PrepareResult | ResultSetGetStreamResponse | None) -> int:
+    """Extract a C ArrowArrayStream pointer from an execute result.
+
+    The pointer is stored as an 8-byte little-endian value inside
+    ``result.stream.value``.  This function validates every step of
+    the extraction and raises descriptive errors on failure, so callers
+    can safely pass the returned integer to Arrow C Data Interface
+    consumers (e.g. PyArrow ``RecordBatchReader.from_stream``).
+
+    Args:
+        result: The protobuf response returned by ``statement_execute``.
+            Must not be ``None`` and must carry a populated ``stream``
+            field.
+
+    Returns:
+        A non-zero integer representing the memory address of the
+        ``ArrowArrayStream`` struct.
+
+    Raises:
+        RuntimeError: If *result* is ``None``, the stream or its value
+            is missing, the value has an unexpected length, or the
+            decoded pointer is null.
+    """
+    if result is None:
+        raise RuntimeError("No query has been executed")
+
+    if not hasattr(result, "stream") or result.stream is None:
+        raise RuntimeError("Execute result does not contain a valid stream")
+
+    if not hasattr(result.stream, "value") or result.stream.value is None:
+        raise RuntimeError("Stream does not contain a valid pointer value")
+
+    stream_value = result.stream.value
+    # 8 bytes = 64-bit pointer, the size of a C ArrowArrayStream* on a 64-bit platform
+    if len(stream_value) != 8:
+        raise RuntimeError(f"Stream pointer value has wrong length: {len(stream_value)} (expected 8)")
+
+    stream_ptr = int.from_bytes(stream_value, byteorder="little", signed=False)
+
+    if stream_ptr == 0:
+        raise RuntimeError("Stream pointer is null")
+
+    return stream_ptr
+
+
+def extract_sqlstate(result: PrepareResult | ResultSetDescriptor | None) -> str | None:
+    """Return the SQLSTATE code from an execute result, if meaningful.
+
+    SQLSTATE ``"00000"`` (successful completion) is normalized to ``None``
+    for backwards compatibility with the legacy connector, which omits
+    the code on success.
+
+    Args:
+        result: A ``PrepareResult`` or ``ResultSetDescriptor``,
+            or ``None`` if no result is available.
+
+    Returns:
+        A five-character SQLSTATE string for warnings/errors, or ``None``
+        on success or when *result* is ``None``.
+    """
+    sql_state = result.sql_state if result else None
+    if sql_state and sql_state != SQLSTATE_SUCCESS:
+        return sql_state
+    return None
+
+
+# Snowflake statement_type_id families/classes that produce a browsable result
+# set or a real update count. When the server omits rows_affected for one of
+# these, the row count is genuinely unknown (-1); for every other classified
+# statement (DDL, session/transaction control, etc.) the legacy connector
+# reported the generic success marker 1, which we reproduce for compatibility.
+# Classification mirrors the core's mask-based taxonomy (level-3 0xF000 family,
+# level-2 0xFF00 class) rather than exact ids, because the concrete ids the
+# server returns (e.g. ALTER SESSION 0x4100, COMMIT 0x5100) do not match named
+# constants but do classify correctly by family.
+_CURSOR_OR_DML_FAMILIES = frozenset(
+    {
+        0x1000,  # SELECT
+        0x2000,  # EXPLAIN
+        0x3000,  # DML (INSERT/UPDATE/DELETE/MERGE/COPY/...)
+        0x7000,  # stage file operations (PUT/GET/LIST/REMOVE)
+        0x9000,  # CALL
+    }
+)
+_CURSOR_CLASSES = frozenset(
+    {
+        0x4400,  # SHOW
+        0x4500,  # DESCRIBE
+        0x4700,  # LIST_FILES
+    }
+)
+_CURSOR_EXACT_IDS = frozenset(
+    {
+        0x6244,  # MANAGE_PATS: DDL-family id that returns a browsable result set
+    }
+)
+# Level-3 (0xF000) families that produce no affected-row count but which the
+# legacy connector reported as rowcount == 1. Restricting the compat fallback
+# to these known families means an unrecognized/future id (e.g. 0xBEEF) is
+# reported as unknown (-1) rather than a spurious success.
+_NO_RESULT_FAMILIES = frozenset(
+    {
+        0x4000,  # SYSCMD (ALTER SESSION 0x4100, USE 0x4300, ...)
+        0x5000,  # TCL (BEGIN/COMMIT/ROLLBACK, e.g. 0x5100, 0x5400)
+        0x6000,  # DDL (CREATE/DROP/ALTER, e.g. 0x6100/0x6101/0x6300)
+        0x8000,  # MISC_QUERY_TYPES
+    }
+)
+
+
+def _produces_result_or_count(statement_type_id: int) -> bool:
+    """Whether the statement type produces a cursor or a real update count.
+
+    These are the statements for which an absent ``rows_affected`` means the
+    count is genuinely unknown; everything else classified is a no-result
+    statement that the legacy connector reported as ``1``.
+    """
+    if (statement_type_id & 0xF000) in _CURSOR_OR_DML_FAMILIES:
+        return True
+    if (statement_type_id & 0xFF00) in _CURSOR_CLASSES:
+        return True
+    return statement_type_id in _CURSOR_EXACT_IDS
+
+
+def extract_rowcount(descriptor: ResultSetDescriptor | None) -> int:
+    """Return the number of rows affected from a ResultSetDescriptor.
+
+    Returns the server's ``rows_affected`` when present (SELECT/DML, including
+    ``0``). When it is absent the statement produced no affected-row count: for
+    a recognized no-result statement (DDL, session/transaction control, etc.)
+    this returns ``1`` to match the legacy connector, and ``-1`` for
+    cursor/DML statements with a missing count or an absent/unrecognized
+    statement type.
+
+    Args:
+        descriptor: The ResultSetDescriptor from a proto response.
+
+    Returns:
+        Row count from server, ``1`` for legacy no-result successes, or ``-1``
+        when unavailable.
+    """
+    if not descriptor:
+        return -1
+
+    if descriptor.HasField("rows_affected"):
+        return descriptor.rows_affected
+
+    if not descriptor.HasField("statement_type_id"):
+        return -1
+    statement_type_id = descriptor.statement_type_id
+    if _produces_result_or_count(statement_type_id):
+        return -1
+    if (statement_type_id & 0xF000) in _NO_RESULT_FAMILIES:
+        return 1
+    return -1