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

snowflake/connector/_internal/cursor/cursor_types.py

@@ -0,0 +1,13 @@
+"""Shared type aliases and type variables for sync and async cursors."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import Any, TypeVar
+
+
+Row = tuple[Any, ...]
+DictRow = dict[str, Any]
+
+F = TypeVar("F", bound=Callable[..., Any])
+Args = TypeVar("Args", bound=Sequence[Any])

snowflake/connector/_internal/cursor/decorators.py

@@ -0,0 +1,56 @@
+"""Cursor precondition decorators shared by sync and async cursor implementations."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ...errors import InterfaceError
+from ..decorators import wrap_method_with_awaitable_pre, wrap_method_with_sync_pre
+from ..errorcode import ER_CURSOR_IS_CLOSED
+from .cursor_types import F
+
+
+def _raise_if_cursor_closed(self: Any) -> None:
+    if self.is_closed():
+        raise InterfaceError(msg="Cursor is closed.", errno=ER_CURSOR_IS_CLOSED)
+
+
+def _raise_if_cursor_flag_closed(self: Any) -> None:
+    if self._closed:
+        raise InterfaceError(msg="Cursor is closed.", errno=ER_CURSOR_IS_CLOSED)
+
+
+def _run_prefetch_hook(self: Any) -> None:
+    if self._prefetch_hook is not None:
+        self._prefetch_hook()
+
+
+async def _run_prefetch_hook_async(self: Any) -> None:
+    if self._prefetch_hook is not None:
+        await self._prefetch_hook()
+
+
+def requires_open(func: F) -> F:
+    """Reject the call when the cursor or its connection is closed."""
+    return wrap_method_with_sync_pre(func, pre=_raise_if_cursor_closed)
+
+
+def requires_open_cursor_not_connection(func: F) -> F:
+    """Guard that only checks ``self._closed``, ignoring the connection state.
+
+    Unlike ``requires_open`` (which delegates to ``is_closed()`` and therefore
+    also rejects cursors whose *connection* has been closed), this decorator
+    deliberately skips the connection check.  This preserves backward
+    compatibility with the old driver, where fetch methods on a cursor with
+    already-buffered results still worked after ``connection.close()``.
+    """
+    return wrap_method_with_sync_pre(func, pre=_raise_if_cursor_flag_closed)
+
+
+def with_prefetch_hook(func: F) -> F:
+    """Invoke the cursor's prefetch hook (if set) before entering the wrapped method."""
+    return wrap_method_with_awaitable_pre(
+        func,
+        sync_pre=_run_prefetch_hook,
+        async_pre=_run_prefetch_hook_async,
+    )

snowflake/connector/_internal/cursor/query_result.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from ...errors import ProgrammingError
+from ..arrow_stream_utils import release_arrow_stream
+from ..protobuf_gen.database_driver_v1_pb2 import (
+    MultiStatementResult,
+    PrepareResult,
+    ResultSetResponse,
+)
+from ..statement_utils import (
+    extract_rowcount,
+    extract_sqlstate,
+    get_stream_ptr,
+)
+from .result_metadata import QueryResultStats, ResultMetadata, ResultMetadataV2
+
+
+class MultiStatementQueryResultState:
+    """Cursor-level navigation state for multi-statement query results.
+
+    Tracks the child query IDs returned by the server, which child
+    we're currently positioned on, and the parent query ID that
+    groups them together.
+
+    Also holds the submitting ``execute()``'s ``request_id``, which the per-child
+    ``ResultSetResponse`` does not carry: this object is the only state
+    ``nextset()`` rescues across ``reset()``, so it is the only place a
+    submission-scoped value can survive to be re-applied to each child.
+    """
+
+    __slots__ = ("parent_qid", "child_query_ids", "request_id", "_next_index")
+
+    def __init__(
+        self,
+        parent_qid: str | None,
+        child_query_ids: list[str],
+        request_id: str | None,
+    ) -> None:
+        self.parent_qid = parent_qid
+        self.child_query_ids = child_query_ids
+        self.request_id = request_id
+        self._next_index = 0
+
+    def advance(self) -> str | None:
+        """Return the next child query ID and advance the index, or None if exhausted."""
+        if self._next_index >= len(self.child_query_ids):
+            return None
+        qid = self.child_query_ids[self._next_index]
+        self._next_index += 1
+        return qid
+
+    def current_child_query_id(self) -> str | None:
+        """Return the query ID of the currently active child result set."""
+        if self._next_index == 0:
+            return None
+        return self.child_query_ids[self._next_index - 1]
+
+    @staticmethod
+    def from_result(
+        multi_result: MultiStatementResult, request_id: str | None
+    ) -> MultiStatementQueryResultState | None:
+        """Create from a proto MultiStatementResult, or None if there are no children.
+
+        ``request_id`` is required: every caller has the submission's requestId
+        in scope, so a None default can never be correct.
+        """
+        query_ids = list(multi_result.query_ids)
+        if not query_ids:
+            return None
+        parent = multi_result.parent
+        parent_qid = parent.query_id if parent.query_id else None
+        return MultiStatementQueryResultState(
+            parent_qid=parent_qid,
+            child_query_ids=query_ids,
+            request_id=request_id,
+        )
+
+
+class QueryResult:
+    """Pure metadata about a query execution result."""
+
+    __slots__ = (
+        "description_v2",
+        "_description_v1",
+        "sqlstate",
+        "sfqid",
+        "request_id",
+        "query",
+        "stats",
+        "rowcount",
+        "is_file_transfer",
+        "query_result_format",
+    )
+
+    def __init__(
+        self,
+        *,
+        description_v2: list[ResultMetadataV2] | None = None,
+        sqlstate: str | None = None,
+        sfqid: str | None = None,
+        request_id: str | None = None,
+        query: str | None = None,
+        stats: QueryResultStats | None = None,
+        rowcount: int | None = None,
+        is_file_transfer: bool = False,
+        query_result_format: str | None = None,
+    ) -> None:
+        self.description_v2 = description_v2
+        self._description_v1: list[ResultMetadata] | None = None
+        self.sqlstate = sqlstate
+        self.sfqid = sfqid
+        # Client-generated requestId of the query submission, populated on
+        # both the success and error paths.
+        self.request_id = request_id
+        self.query = query
+        self.stats = stats if stats is not None else QueryResultStats()
+        self.rowcount = rowcount
+        self.is_file_transfer = is_file_transfer
+        self.query_result_format = query_result_format
+
+    @property
+    def description(self) -> list[ResultMetadata] | None:
+        """PEP 249 view of the column metadata, downcast from ``description_v2``.
+
+        ``description_v2`` is the single parse of the proto column list; this
+        view is built on first read and cached.
+        """
+        if self._description_v1 is None and self.description_v2 is not None:
+            self._description_v1 = [col._to_result_metadata_v1() for col in self.description_v2]
+        return self._description_v1
+
+    def reset(self, closing: bool = False) -> None:
+        """Optionally clear the rowcount.
+
+        Only rowcount is reset — description, sqlstate, sfqid, request_id, query,
+        and stats are left intact for backward compatibility (callers may read them after close/reset).
+        They are overwritten wholesale when the cursor's _query_result is replaced on the next execute().
+        """
+        if not closing:
+            self.rowcount = None
+
+    @staticmethod
+    def from_prepare_result(result: PrepareResult | None) -> QueryResult:
+        stream_ptr = get_stream_ptr(result)
+        release_arrow_stream(stream_ptr)
+
+        description = ResultMetadataV2.create_description(result)
+        return QueryResult(
+            description_v2=description,
+            sqlstate=extract_sqlstate(result),
+            sfqid=(result.query_id if result.query_id else None) if result else None,
+            request_id=(result.request_id if result.request_id else None) if result else None,
+            query=(result.query if result.query else None) if result else None,
+            rowcount=0 if description else None,
+        )
+
+    @staticmethod
+    def from_programming_error(exc: ProgrammingError) -> QueryResult:
+        return QueryResult(
+            sqlstate=exc.sqlstate or None,
+            sfqid=exc.sfqid or None,
+            request_id=exc.request_id or None,
+            query=exc.query or None,
+        )
+
+    @staticmethod
+    def from_result_set_response(
+        response: ResultSetResponse,
+        query: str | None = None,
+        request_id: str | None = None,
+    ) -> QueryResult:
+        """Create QueryResult from a ResultSetResponse (metadata only).
+
+        Args:
+            response: ResultSetResponse containing descriptor metadata.
+            query: Optional query text (not available in proto, must be passed separately).
+            request_id: Optional submission requestId (not carried on ResultSetResponse,
+                must be passed separately by callers that submitted the query).
+
+        Returns:
+            QueryResult instance with metadata populated.
+        """
+        descriptor = response.result_descriptor
+
+        return QueryResult(
+            description_v2=ResultMetadataV2.create_description(descriptor),
+            sqlstate=extract_sqlstate(descriptor),
+            sfqid=descriptor.query_id if descriptor.query_id else None,
+            request_id=request_id,
+            query=query,
+            rowcount=extract_rowcount(descriptor),
+            is_file_transfer=(
+                descriptor.statement_type_id in (0x7101, 0x7102) if descriptor.HasField("statement_type_id") else False
+            ),
+            stats=(
+                QueryResultStats.from_query_stats(descriptor.stats)
+                if descriptor.HasField("stats")
+                else QueryResultStats()
+            ),
+            # Legacy reports a format-less response as JSON.
+            query_result_format=descriptor.query_result_format or "json",
+        )

snowflake/connector/_internal/cursor/query_result_waiter.py

@@ -0,0 +1,98 @@
+"""Waiter that polls query status until a query completes or errors."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+
+from typing import TYPE_CHECKING
+
+from ...constants import QueryStatus
+from ...errors import DatabaseError
+
+
+if TYPE_CHECKING:
+    from ...aio.connection import Connection as AsyncConnection
+    from ...connection import Connection
+    from ..connection import ConnectionMixin
+
+_RETRY_PATTERN = [1, 1, 2, 3, 4, 8, 10]
+_NO_DATA_MAX_RETRY = 24
+
+
+def _no_data_error(sfqid: str) -> DatabaseError:
+    return DatabaseError(
+        f"Cannot retrieve data on the status of this query. No information returned from server for query '{sfqid}'"
+    )
+
+
+def _advance_no_data_counter(status: QueryStatus, no_data_counter: int, sfqid: str) -> int:
+    if status != QueryStatus.NO_DATA:
+        return no_data_counter
+    no_data_counter += 1
+    if no_data_counter > _NO_DATA_MAX_RETRY:
+        raise _no_data_error(sfqid)
+    return no_data_counter
+
+
+def _poll_step(
+    connection: ConnectionMixin,
+    sfqid: str,
+    status: QueryStatus,
+    no_data_counter: int,
+    retry_idx: int,
+) -> tuple[float | None, int, int]:
+    """Advance one poll iteration.
+
+    Returns ``(sleep_seconds, no_data_counter, retry_idx)``. ``sleep_seconds``
+    is ``None`` when polling should stop.
+    """
+    if not connection.is_still_running(status):
+        return None, no_data_counter, retry_idx
+    no_data_counter = _advance_no_data_counter(status, no_data_counter, sfqid)
+    delay = 0.5 * _RETRY_PATTERN[retry_idx]
+    if retry_idx < len(_RETRY_PATTERN) - 1:
+        retry_idx += 1
+    return delay, no_data_counter, retry_idx
+
+
+class QueryResultWaiter:
+    """Sync waiter that polls query status with capped exponential backoff."""
+
+    def __init__(self, connection: Connection, sfqid: str) -> None:
+        self._connection = connection
+        self._sfqid = sfqid
+
+    def wait(self) -> None:
+        """Block until the query completes. Raise on terminal error status."""
+        no_data_counter = 0
+        retry_idx = 0
+        while True:
+            status = self._connection.get_query_status_throw_if_error(self._sfqid)
+            delay, no_data_counter, retry_idx = _poll_step(
+                self._connection, self._sfqid, status, no_data_counter, retry_idx
+            )
+            if delay is None:
+                return
+            time.sleep(delay)
+
+
+class AsyncQueryResultWaiter:
+    """Async waiter that polls query status with capped exponential backoff."""
+
+    def __init__(self, connection: AsyncConnection, sfqid: str) -> None:
+        self._connection = connection
+        self._sfqid = sfqid
+
+    async def wait(self) -> None:
+        """Await query completion. Raise on terminal error status."""
+        no_data_counter = 0
+        retry_idx = 0
+        while True:
+            status = await self._connection.get_query_status_throw_if_error(self._sfqid)
+            delay, no_data_counter, retry_idx = _poll_step(
+                self._connection, self._sfqid, status, no_data_counter, retry_idx
+            )
+            if delay is None:
+                return
+            await asyncio.sleep(delay)

snowflake/connector/_internal/cursor/result_metadata.py

@@ -0,0 +1,255 @@
+"""Result metadata and query statistics types."""
+
+from __future__ import annotations
+
+from typing import Any, NamedTuple
+
+from ..protobuf_gen.database_driver_v1_pb2 import (
+    PrepareResult,
+    QueryStats,
+    ResultSetDescriptor,
+)
+from ..type_codes import get_type_code
+
+
+_TEXT_TYPES = ("TEXT", "VARCHAR", "CHAR", "STRING")
+
+
+def _type_code_name(col: Any) -> str:
+    ext = getattr(col, "ext_col_type_name", "")
+    if isinstance(ext, str) and ext:
+        return ext
+    return str(col.type)
+
+
+def _column_display_size(col: Any) -> int | None:
+    # Char count, text-only — a new-driver-only enhancement (BD#90); the old driver leaves it unset.
+    return col.length if col.HasField("length") and col.type.upper() in _TEXT_TYPES else None
+
+
+def _column_internal_size(col: Any) -> int | None:
+    # Char count (old-driver semantics), ungated — BINARY/VARBINARY need a real value too.
+    return col.length if col.HasField("length") else None
+
+
+class ResultMetadata(NamedTuple):
+    """PEP 249 column description entry.
+
+    Each item in ``Cursor.description`` is a ``ResultMetadata`` instance.
+    Being a :class:`~typing.NamedTuple` it is fully tuple-compatible as
+    required by the spec, while also providing named attribute access.
+    """
+
+    name: str
+    type_code: int
+    display_size: int | None
+    internal_size: int | None
+    precision: int | None
+    scale: int | None
+    is_nullable: bool | None
+
+    @classmethod
+    def from_column(cls, col: Any) -> ResultMetadata:
+        """Create a ``ResultMetadata`` from a protobuf ``ColumnMetadata``."""
+        type_code = get_type_code(_type_code_name(col))
+
+        display_size = _column_display_size(col)
+        internal_size = _column_internal_size(col)
+        precision = col.precision if col.HasField("precision") else None
+        scale = col.scale if col.HasField("scale") else None
+
+        return cls(
+            name=col.name,
+            type_code=type_code,
+            display_size=display_size,
+            internal_size=internal_size,
+            precision=precision,
+            scale=scale,
+            is_nullable=col.nullable,
+        )
+
+
+class ResultMetadataV2:
+    """New-format column description carrying ``vector_dimension`` and ``fields``.
+
+    Matches the legacy ``snowflake-connector-python`` ``ResultMetadataV2``
+    interface so Snowpark can read ``.name``, ``.type_code``, ``.is_nullable``,
+    ``.vector_dimension``, and ``.fields`` without connector-version guards.
+
+    ``vector_dimension`` is populated from the proto ``dimension`` field.
+    ``fields`` carries nested metadata for structured types: the element type
+    for VECTOR and ARRAY, the key and value types for MAP, and one entry per
+    attribute for structured OBJECT. It is ``None`` for every other type, and
+    also for semi-structured ARRAY/OBJECT/MAP, which the server describes
+    without a nested field list.
+
+    Note: ``_is_nullable`` is exposed both as the ``is_nullable`` public property
+    and as the ``_is_nullable`` private attribute, because Snowpark accesses the
+    private attribute directly on nested ARRAY/MAP element metadata objects
+    (``obj._is_nullable``). Both must exist by their exact names.
+    """
+
+    __slots__ = (
+        "_name",
+        "_type_code",
+        "_is_nullable",
+        "_display_size",
+        "_internal_size",
+        "_precision",
+        "_scale",
+        "_vector_dimension",
+        "_fields",
+    )
+
+    def __init__(
+        self,
+        name: str | None,
+        type_code: int,
+        is_nullable: bool,
+        display_size: int | None = None,
+        internal_size: int | None = None,
+        precision: int | None = None,
+        scale: int | None = None,
+        vector_dimension: int | None = None,
+        fields: list[ResultMetadataV2] | None = None,
+    ) -> None:
+        self._name = name
+        self._type_code = type_code
+        self._is_nullable = is_nullable
+        self._display_size = display_size
+        self._internal_size = internal_size
+        self._precision = precision
+        self._scale = scale
+        self._vector_dimension = vector_dimension
+        self._fields = fields
+
+    @property
+    def name(self) -> str | None:
+        return self._name
+
+    @property
+    def type_code(self) -> int:
+        return self._type_code
+
+    @property
+    def is_nullable(self) -> bool:
+        return self._is_nullable
+
+    @property
+    def display_size(self) -> int | None:
+        return self._display_size
+
+    @property
+    def internal_size(self) -> int | None:
+        return self._internal_size
+
+    @property
+    def precision(self) -> int | None:
+        return self._precision
+
+    @property
+    def scale(self) -> int | None:
+        return self._scale
+
+    @property
+    def vector_dimension(self) -> int | None:
+        return self._vector_dimension
+
+    @property
+    def fields(self) -> list[ResultMetadataV2] | None:
+        return self._fields
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ResultMetadataV2):
+            return NotImplemented
+        return (
+            self._name == other._name
+            and self._type_code == other._type_code
+            and self._is_nullable == other._is_nullable
+            and self._display_size == other._display_size
+            and self._internal_size == other._internal_size
+            and self._precision == other._precision
+            and self._scale == other._scale
+            and self._vector_dimension == other._vector_dimension
+            and self._fields == other._fields
+        )
+
+    def _to_result_metadata_v1(self) -> ResultMetadata:
+        """Downcast to a PEP 249-compatible ``ResultMetadata`` NamedTuple."""
+        return ResultMetadata(
+            name=self._name,  # type: ignore[arg-type]  # sub-fields carry name=None (legacy parity)
+            type_code=self._type_code,
+            display_size=self._display_size,
+            internal_size=self._internal_size,
+            precision=self._precision,
+            scale=self._scale,
+            is_nullable=self._is_nullable,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"ResultMetadataV2(name={self._name!r}, type_code={self._type_code!r}, "
+            f"is_nullable={self._is_nullable!r}, vector_dimension={self._vector_dimension!r}, "
+            f"fields={self._fields!r})"
+        )
+
+    @classmethod
+    def from_column(cls, col: Any, *, nested: bool = False) -> ResultMetadataV2:
+        """Build from a proto ``ColumnMetadata`` message.
+
+        ``nested`` marks a sub-field of a structured type. Those carry a name
+        only for the attributes of a structured OBJECT; the proto renders an
+        absent name as the empty string, which maps back to ``None`` here for
+        parity with the old driver.
+        """
+        type_code = get_type_code(_type_code_name(col))
+        display_size = _column_display_size(col)
+        internal_size = _column_internal_size(col)
+        precision = col.precision if col.HasField("precision") else None
+        scale = col.scale if col.HasField("scale") else None
+        vector_dimension = col.dimension if col.HasField("dimension") else None
+        return cls(
+            name=(col.name or None) if nested else col.name,
+            type_code=type_code,
+            is_nullable=col.nullable,
+            display_size=display_size,
+            internal_size=internal_size,
+            precision=precision,
+            scale=scale,
+            vector_dimension=vector_dimension,
+            # For MAP, `col.fields` carries exactly two entries, key type then value
+            # type, in that order. Snowpark's structured-type inference reads them
+            # positionally as fields[0]/fields[1] today, so this recursion must
+            # preserve wire order rather than resort or drop an entry.
+            fields=[cls.from_column(f, nested=True) for f in col.fields] if col.fields else None,
+        )
+
+    @classmethod
+    def create_description(cls, result: PrepareResult | ResultSetDescriptor | None) -> list[ResultMetadataV2] | None:
+        """Build a V2 description list from a prepare/describe result."""
+        if result and result.columns:
+            return [cls.from_column(col) for col in result.columns]
+        return None
+
+
+class QueryResultStats(NamedTuple):
+    """DML operation statistics returned by Snowflake.
+
+    Exposes per-operation row counts for INSERT, UPDATE, DELETE,
+    and the number of duplicate rows skipped during DML.
+    """
+
+    num_rows_inserted: int | None = None
+    num_rows_deleted: int | None = None
+    num_rows_updated: int | None = None
+    num_dml_duplicates: int | None = None
+
+    @classmethod
+    def from_query_stats(cls, s: QueryStats) -> QueryResultStats:
+        """Create a ``QueryResultStats`` from a protobuf ``QueryStats``."""
+        return cls(
+            num_rows_inserted=s.num_rows_inserted if s.HasField("num_rows_inserted") else None,
+            num_rows_deleted=s.num_rows_deleted if s.HasField("num_rows_deleted") else None,
+            num_rows_updated=s.num_rows_updated if s.HasField("num_rows_updated") else None,
+            num_dml_duplicates=s.num_dml_duplicates if s.HasField("num_dml_duplicates") else None,
+        )

snowflake/connector/_internal/cursor/result_set_wrapper.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from ...errors import ProgrammingError
+from ..api_client.client_api import core_driver
+from ..errorcode import ER_NO_DATA_FOUND
+from ..logging import get_logger
+from ..protobuf_gen.database_driver_v1_pb2 import ResultSetHandle
+
+
+logger = get_logger(__name__)
+
+
+class ResultSetWrapperBase:
+    """Owns a ``ResultSetHandle`` and guards against misuse.
+
+    Lifecycle: ``replace()`` releases the previous handle before adopting a
+    new one; ``release()`` frees the current handle.  ``__del__`` acts as a
+    safety net for handles that were never explicitly released.
+
+    The core layer supports repeated ``result_set_get_stream`` calls — each
+    invocation builds a fresh Arrow stream from the stored ``RowsetData``.
+
+    Release/replace are always synchronous (pure FFI handle deletion, no network I/O).
+    """
+
+    __slots__ = ("_handle",)
+
+    def __init__(self, handle: ResultSetHandle | None = None) -> None:
+        self._handle: ResultSetHandle | None = handle
+
+    def __del__(self) -> None:
+        self._do_release()
+
+    def replace(self, new_handle: ResultSetHandle | None) -> None:
+        """Release the current handle (if any) and adopt *new_handle*."""
+        self._do_release()
+        self._handle = new_handle
+
+    def release(self) -> None:
+        """Explicitly release the handle."""
+        self._do_release()
+
+    def _do_release(self) -> None:
+        handle = self._take_handle()
+        if handle is None:
+            return
+        try:
+            core_driver.result_set_release(result_set_handle=handle)
+        except Exception as e:
+            logger.warning("Failed to release ResultSet handle: %s", type(e).__name__)
+            logger.debug("Failed to release ResultSet handle", exc_info=True)
+
+    def _take_handle(self) -> ResultSetHandle | None:
+        handle = self._handle
+        if handle is None:
+            return None
+        self._handle = None
+        return handle
+
+    def _require_handle(self) -> ResultSetHandle:
+        if self._handle is None:
+            raise ProgrammingError(
+                msg="No results available (not produced by this query)",
+                errno=ER_NO_DATA_FOUND,
+            )
+        return self._handle