snowflake-sql-api 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

snowflake_sql_api/__init__.py

@@ -27,8 +27,10 @@ except ImportError:  # pragma: no cover - plain source checkout, never built
         __version__ = "0.0.0.dev0"
 
 from .aclient import AsyncQueryHandle, AsyncSnowflakeClient
+from .batches import BatchSummary, BatchWriteResult, FailedBatch
 from .client import QueryHandle, SnowflakeClient
 from .exceptions import (
+    BatchWriteError,
     ResultNotReady,
     SnowflakeAuthError,
     SnowflakeConfigError,
@@ -48,6 +50,10 @@ __all__ = [
     "QueryHandle",
     "AsyncQueryHandle",
     "RetryPolicy",
+    "BatchSummary",
+    "BatchWriteResult",
+    "FailedBatch",
+    "BatchWriteError",
     "SnowflakeError",
     "SnowflakeConfigError",
     "SnowflakeAuthError",

snowflake_sql_api/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.1.2'
-__version_tuple__ = version_tuple = (0, 1, 2)
+__version__ = version = '0.1.4'
+__version_tuple__ = version_tuple = (0, 1, 4)
 
 __commit_id__ = commit_id = None

snowflake_sql_api/aclient.py

@@ -12,11 +12,23 @@ import os
 import time
 import uuid
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import (
+    Any,
+    AsyncGenerator,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    TextIO,
+    Tuple,
+)
 
 import httpx
 
 from .auth import KeypairAuthenticator, account_hostname
+from .batches import BatchWriteResult, _BatchExecution, _BatchPlanner
 from .bindings import to_bindings
 from .client import DEFAULT_POLL_TIMEOUT, _columns_from, _rows_affected
 from .escaping import quote_identifier, quote_name
@@ -26,7 +38,13 @@ from .exceptions import (
     SnowflakeRequestError,
     SnowflakeTimeoutError,
 )
-from .pagination import fetch_all_partitions_async
+from .exports import (
+    CSVWriter,
+    validate_batch_size,
+    validate_export_columns,
+    write_jsonl,
+)
+from .pagination import fetch_all_partitions_async, iter_result_rows_async
 from .retry import DEFAULT_RETRY_POLICY, RetryPolicy
 from .transport import AsyncTransport, StatementResponse
 from .types import ColumnMeta, Row, coerce_rows
@@ -201,13 +219,13 @@ class AsyncSnowflakeClient:
                 )
             await asyncio.sleep(self._poll_interval)
 
-    async def _collect(
+    async def _prepare_result(
         self,
         response: StatementResponse,
         *,
         timeout: Optional[float],
         poll: bool = True,
-    ) -> Tuple[List[ColumnMeta], List[List[Any]]]:
+    ) -> Tuple[List[ColumnMeta], str, Dict[str, Any]]:
         if response.is_running:
             if not poll:
                 raise ResultNotReady(statement_handle=response.statement_handle)
@@ -220,8 +238,27 @@ class AsyncSnowflakeClient:
         handle = response.statement_handle
         if handle is None:
             raise SnowflakeRequestError("completed statement has no handle")
-        raw_rows = await fetch_all_partitions_async(self._transport, handle, body)
-        return columns, raw_rows
+        return columns, handle, body
+
+    async def _collect(
+        self,
+        response: StatementResponse,
+        *,
+        timeout: Optional[float],
+        poll: bool = True,
+    ) -> Tuple[List[ColumnMeta], List[List[Any]]]:
+        columns, handle, body = await self._prepare_result(
+            response, timeout=timeout, poll=poll
+        )
+        return columns, await fetch_all_partitions_async(self._transport, handle, body)
+
+    async def _stream_result(
+        self, sql: str, params: Params
+    ) -> Tuple[List[ColumnMeta], AsyncGenerator[Row, None]]:
+        columns, handle, body = await self._prepare_result(
+            await self._run(sql, params), timeout=self._statement_timeout
+        )
+        return columns, iter_result_rows_async(self._transport, handle, body, columns)
 
     # -- query surface ----------------------------------------------------
 
@@ -233,6 +270,73 @@ class AsyncSnowflakeClient:
         )
         return coerce_rows(raw_rows, columns)
 
+    async def query_stream(
+        self, sql: str, params: Params = None
+    ) -> AsyncGenerator[Row, None]:
+        """Yield converted rows on demand; call directly with async for.
+
+        Submission starts on first iteration. Keep the client open while
+        consuming; await aclose() on a retained stream when stopping early.
+        """
+        _, rows = await self._stream_result(sql, params)
+        try:
+            async for row in rows:
+                yield row
+        finally:
+            await rows.aclose()
+
+    async def query_batches(
+        self, sql: str, params: Params = None, *, batch_size: int = 1000
+    ) -> AsyncGenerator[List[Row], None]:
+        """Yield batches of at most batch_size rows across partition boundaries."""
+        validate_batch_size(batch_size)
+        rows = self.query_stream(sql, params)
+        try:
+            batch: List[Row] = []
+            async for row in rows:
+                batch.append(row)
+                if len(batch) == batch_size:
+                    yield batch
+                    batch = []
+            if batch:
+                yield batch
+        finally:
+            await rows.aclose()
+
+    async def export_csv(self, sql: str, output: TextIO, params: Params = None) -> int:
+        """Await query I/O and write CSV to a caller-owned synchronous text sink.
+
+        Return the data-row count. Failures may leave partial output.
+        """
+        columns, rows = await self._stream_result(sql, params)
+        try:
+            writer = CSVWriter(output, columns)
+            count = 0
+            async for row in rows:
+                writer.write(row)
+                count += 1
+        finally:
+            await rows.aclose()
+        return count
+
+    async def export_jsonl(
+        self, sql: str, output: TextIO, params: Params = None
+    ) -> int:
+        """Await query I/O and write JSONL to a caller-owned synchronous text sink.
+
+        Return the row count. Failures may leave partial output.
+        """
+        columns, rows = await self._stream_result(sql, params)
+        try:
+            validate_export_columns(columns)
+            count = 0
+            async for row in rows:
+                write_jsonl(output, row)
+                count += 1
+        finally:
+            await rows.aclose()
+        return count
+
     async def query_one(self, sql: str, params: Params = None) -> Optional[Row]:
         """Return the first row, or ``None`` if the result is empty."""
         rows = await self.query(sql, params)
@@ -286,6 +390,52 @@ class AsyncSnowflakeClient:
 
     # -- async submit -----------------------------------------------------
 
+    async def insert_batches(
+        self,
+        table: str,
+        columns: Sequence[str],
+        rows: Iterable[Sequence[Any]],
+        *,
+        max_rows: int = 1000,
+        max_bytes: int = 1_000_000,
+    ) -> BatchWriteResult:
+        """Insert a single-pass source sequentially within row and wire-byte budgets.
+
+        Independent statements can commit before a later failure. BatchWriteError
+        reports confirmed progress and uncertain outcomes; see docs/batch-writes.md.
+        Input is an ordinary iterable, consumed on the calling thread.
+        """
+        planner = _BatchPlanner(
+            table, columns, rows, max_rows, max_bytes, self._build_payload
+        )
+        while (batch := planner.next_batch()) is not None:
+            try:
+                self._notify(batch.sql, batch.params)
+            except Exception as exc:
+                raise planner.error(exc, "validation") from exc
+            execution = _BatchExecution(planner, batch)
+            try:
+                response = await self._transport.submit(
+                    {}, request_id=execution.failed.request_id, content=batch.content
+                )
+                execution.received(
+                    response.statement_handle,
+                    response.submission_attempts,
+                    response.is_running,
+                )
+                if response.is_running:
+                    handle = response.statement_handle
+                    if handle is None:
+                        raise SnowflakeRequestError("running statement has no handle")
+                    response = await self._poll(handle, self._statement_timeout)
+                    execution.succeeded = True
+                affected = _rows_affected(response.body, strict=True)
+            except Exception as exc:
+                raise execution.error(exc) from exc
+            planner.confirm(batch, affected)
+            del batch
+        return planner.result()
+
     async def submit(self, sql: str, params: Params = None) -> AsyncQueryHandle:
         """Submit a long-running statement and return an :class:`AsyncQueryHandle`."""
         response = await self._run(sql, params, async_exec=True)

snowflake_sql_api/batches.py

@@ -0,0 +1,298 @@
+"""Bounded INSERT planning and immutable progress shared by both clients."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import Iterable, Iterator, Sequence
+from dataclasses import dataclass, replace
+from typing import Any, Callable, Literal
+
+from .bindings import to_bindings
+from .escaping import quote_identifier, quote_name
+from .exceptions import (
+    BatchWriteError,
+    SnowflakeAuthError,
+    SnowflakeConfigError,
+    SnowflakeProgrammingError,
+    SnowflakeRequestError,
+)
+
+__all__ = ["BatchSummary", "BatchWriteResult", "FailedBatch"]
+
+BatchPhase = Literal["input", "validation", "submission", "execution"]
+BatchOutcome = Literal["not_executed", "failed", "unknown", "succeeded"]
+
+
+@dataclass(frozen=True)
+class BatchSummary:
+    """One confirmed INSERT, with zero-based batch and source offsets."""
+
+    index: int
+    start_row_index: int
+    row_count: int
+    rows_affected: int
+
+
+@dataclass(frozen=True)
+class BatchWriteResult:
+    """Confirmed progress only; never retains input records or request bodies."""
+
+    batches: tuple[BatchSummary, ...] = ()
+
+    @property
+    def input_rows(self) -> int:
+        return sum(batch.row_count for batch in self.batches)
+
+    @property
+    def rows_affected(self) -> int:
+        return sum(batch.rows_affected for batch in self.batches)
+
+    @property
+    def completed_batches(self) -> int:
+        return len(self.batches)
+
+
+@dataclass(frozen=True)
+class FailedBatch:
+    """Prepared INSERT identity; an unknown outcome requires reconciliation."""
+
+    index: int
+    start_row_index: int
+    row_count: int
+    request_id: str
+    statement_handle: str | None = None
+    submission_attempts: int = 0
+
+
+def _encode(value: Any) -> bytes:
+    return json.dumps(
+        value, ensure_ascii=False, separators=(",", ":"), allow_nan=False
+    ).encode("utf-8")
+
+
+@dataclass
+class _PreparedBatch:
+    sql: str
+    params: list[Any]
+    content: bytes
+    row_count: int
+
+
+class _BatchPlanner:
+    def __init__(
+        self,
+        table: str,
+        columns: Sequence[str],
+        rows: Iterable[Sequence[Any]],
+        max_rows: int,
+        max_bytes: int,
+        build_payload: Callable[[str, Any], dict[str, Any]],
+    ) -> None:
+        if not columns:
+            raise SnowflakeConfigError("insert_batches requires at least one column")
+        for name, limit in (("max_rows", max_rows), ("max_bytes", max_bytes)):
+            if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0:
+                raise SnowflakeConfigError(f"{name} must be a positive integer")
+        if max_rows > 200_000:
+            raise SnowflakeConfigError("max_rows must be <= 200000")
+        cols = ", ".join(quote_identifier(c) for c in columns)
+        self.prefix = f"INSERT INTO {quote_name(table)} ({cols}) VALUES "
+        self.placeholder = "(" + ", ".join(["?"] * len(columns)) + ")"
+        self.width = len(columns)
+        self.max_rows, self.max_bytes = max_rows, max_bytes
+        self.source = rows
+        self.iterator: Iterator[Sequence[Any]] | None = None
+        self.build_payload = build_payload
+        self.context: bytes | None = None
+        self.single_overhead = 0
+        self.pending: tuple[list[Any], list[bytes]] | None = None
+        self.exhausted = False
+        self.consumed_rows = 0
+        self.confirmed_rows = 0
+        self.completed: list[BatchSummary] = []
+
+    def result(self) -> BatchWriteResult:
+        return BatchWriteResult(tuple(self.completed))
+
+    def error(
+        self,
+        cause: Exception,
+        phase: BatchPhase,
+        outcome: BatchOutcome = "not_executed",
+        failed_batch: FailedBatch | None = None,
+    ) -> BatchWriteError:
+        return BatchWriteError(
+            f"batch write stopped during {phase}: {cause}",
+            completed=self.result(),
+            next_row_index=self.confirmed_rows,
+            consumed_rows=self.consumed_rows,
+            failed_batch=failed_batch,
+            phase=phase,
+            outcome=outcome,
+        )
+
+    def _record(self) -> tuple[list[Any], list[bytes]] | None:
+        if self.exhausted:
+            return None
+        try:
+            if self.iterator is None:
+                self.iterator = iter(self.source)
+        except Exception as exc:
+            raise self.error(exc, "input") from exc
+        try:
+            row = next(self.iterator)
+        except StopIteration:
+            self.exhausted = True
+            return None
+        except Exception as exc:
+            raise self.error(exc, "input") from exc
+        self.consumed_rows += 1
+        try:
+            if not isinstance(row, Sequence) or isinstance(
+                row, (str, bytes, bytearray)
+            ):
+                raise SnowflakeConfigError("each row must be a positional sequence")
+            if len(row) != self.width:
+                raise SnowflakeConfigError(
+                    f"each row must have {self.width} values, got {len(row)}"
+                )
+            values = list(row)
+            return values, [_encode(value) for value in to_bindings(values).values()]
+        except Exception as exc:
+            raise self.error(exc, "validation") from exc
+
+    @staticmethod
+    def _fragment(values: list[bytes], offset: int) -> bytes:
+        return b",".join(
+            _encode(str(offset + i)) + b":" + value for i, value in enumerate(values, 1)
+        )
+
+    def _body(self, sql: str, bindings: bytes) -> bytes:
+        assert self.context is not None
+        return (
+            b'{"statement":'
+            + _encode(sql)
+            + self.context
+            + b',"bindings":{'
+            + bindings
+            + b"}}"
+        )
+
+    def next_batch(self) -> _PreparedBatch | None:
+        # Configuration is validated before this point; input errors discard the
+        # unsent buffer. Only a validated singleton can trigger a byte-boundary flush.
+        params: list[Any] = []
+        fragments: list[bytes] = []
+        row_count = 0
+        size = 0
+        while row_count < self.max_rows:
+            record = self.pending
+            self.pending = None
+            if record is None:
+                record = self._record()
+            if record is None:
+                break
+            values, encoded = record
+            try:
+                if self.context is None:
+                    context = self.build_payload("", None)
+                    del context["statement"]
+                    self.context = _encode(context)[1:-1]
+                    if self.context:
+                        self.context = b"," + self.context
+                    self.single_overhead = len(
+                        self._body(self.prefix + self.placeholder, b"")
+                    )
+                first = self._fragment(encoded, 0)
+                single_size = self.single_overhead + len(first)
+                if single_size > self.max_bytes:
+                    raise SnowflakeConfigError(
+                        f"row {self.consumed_rows - 1} exceeds max_bytes={self.max_bytes}"
+                    )
+                fragment = (
+                    first if row_count == 0 else self._fragment(encoded, len(params))
+                )
+                candidate = (
+                    single_size
+                    if row_count == 0
+                    else size + 2 + len(self.placeholder) + 1 + len(fragment)
+                )
+                if candidate > self.max_bytes:
+                    self.pending = record
+                    break
+                fragments.append(fragment)
+                params.extend(values)
+                row_count += 1
+                size = candidate
+            except Exception as exc:
+                raise self.error(exc, "validation") from exc
+        if not row_count:
+            return None
+        sql = self.prefix + ", ".join([self.placeholder] * row_count)
+        content = self._body(sql, b",".join(fragments))
+        if len(content) != size or len(content) > self.max_bytes:
+            budget_error = SnowflakeConfigError(
+                "batch body exceeds its measured byte budget"
+            )
+            raise self.error(budget_error, "validation") from budget_error
+        return _PreparedBatch(sql, params, content, row_count)
+
+    def confirm(self, batch: _PreparedBatch, rows_affected: int) -> None:
+        self.completed.append(
+            BatchSummary(
+                len(self.completed), self.confirmed_rows, batch.row_count, rows_affected
+            )
+        )
+        self.confirmed_rows += batch.row_count
+
+
+class _BatchExecution:
+    """Parent POST context survives polling and conservative error classification."""
+
+    def __init__(self, planner: _BatchPlanner, batch: _PreparedBatch) -> None:
+        self.planner = planner
+        self.failed = FailedBatch(
+            len(planner.completed),
+            planner.confirmed_rows,
+            batch.row_count,
+            str(uuid.uuid4()),
+        )
+        self.phase: BatchPhase = "submission"
+        self.succeeded = False
+
+    def received(self, handle: str | None, attempts: int | None, running: bool) -> None:
+        self.failed = replace(
+            self.failed, statement_handle=handle, submission_attempts=attempts or 0
+        )
+        if not running or handle is not None:
+            self.phase = "execution"
+        self.succeeded = not running
+
+    def error(self, cause: Exception) -> BatchWriteError:
+        attempts = getattr(cause, "submission_attempts", None)
+        handle = getattr(cause, "statement_handle", None) or getattr(
+            cause.__cause__, "statement_handle", None
+        )
+        self.failed = replace(
+            self.failed,
+            submission_attempts=(
+                self.failed.submission_attempts if attempts is None else attempts
+            ),
+            statement_handle=handle or self.failed.statement_handle,
+        )
+        outcome: BatchOutcome = "succeeded" if self.succeeded else "unknown"
+        if not self.succeeded and self.failed.submission_attempts <= 1:
+            if isinstance(cause, SnowflakeProgrammingError):
+                self.phase, outcome = "execution", "failed"
+            elif self.phase == "submission" and (
+                attempts == 0
+                or isinstance(cause, SnowflakeAuthError)
+                and cause.status_code == 401
+                or isinstance(cause, SnowflakeRequestError)
+                and cause.status_code in (400, 401, 403)
+            ):
+                outcome = "not_executed"
+        if isinstance(cause, SnowflakeProgrammingError):
+            self.phase = "execution"
+        return self.planner.error(cause, self.phase, outcome, self.failed)