snowflake-sql-api 0.1.3__tar.gz → 0.1.4__tar.gz

{snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: snowflake-sql-api
-Version: 0.1.3
+Version: 0.1.4
 Summary: Lightweight, pure-Python client for Snowflake's SQL API v2
 Project-URL: Homepage, https://github.com/hampsterx/snowflake-sql-api
 Project-URL: Repository, https://github.com/hampsterx/snowflake-sql-api
@@ -155,6 +155,25 @@ with the largest decoded partition plus the application batch. See
 [streaming results and exports](docs/streaming.md) for row streams, async usage,
 serialization, and partial-output behavior.
 
+Insert a record stream in bounded requests:
+
+```python
+with SnowflakeClient.from_env() as client:
+    progress = client.insert_batches(
+        "events", ["id", "label"], ((i, f"event {i}") for i in range(2500)),
+        max_rows=1000, max_bytes=1_000_000,
+    )
+    print(progress.input_rows, progress.rows_affected)
+```
+
+`insert_many()` sends one statement; `insert_batches()` sends sequential statements
+and can partially complete. Its byte budget covers the exact JSON body. On failure,
+`BatchWriteError` reports confirmed progress, consumed input and the failed request's
+identity. Reconstruct replayable input and reconcile unknown outcomes before retrying;
+do not resume with the same advanced iterator. The async equivalent is awaited and
+accepts the same synchronous iterable, whose production can block the event loop.
+See [batch writes](docs/batch-writes.md) for limits and recovery semantics.
+
 ## Prior art and alternatives
 
 `snowflake-sql-api` is not the first REST-based Snowflake client. If it does not

{snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.4}/README.md

@@ -91,6 +91,25 @@ with the largest decoded partition plus the application batch. See
 [streaming results and exports](docs/streaming.md) for row streams, async usage,
 serialization, and partial-output behavior.
 
+Insert a record stream in bounded requests:
+
+```python
+with SnowflakeClient.from_env() as client:
+    progress = client.insert_batches(
+        "events", ["id", "label"], ((i, f"event {i}") for i in range(2500)),
+        max_rows=1000, max_bytes=1_000_000,
+    )
+    print(progress.input_rows, progress.rows_affected)
+```
+
+`insert_many()` sends one statement; `insert_batches()` sends sequential statements
+and can partially complete. Its byte budget covers the exact JSON body. On failure,
+`BatchWriteError` reports confirmed progress, consumed input and the failed request's
+identity. Reconstruct replayable input and reconcile unknown outcomes before retrying;
+do not resume with the same advanced iterator. The async equivalent is awaited and
+accepts the same synchronous iterable, whose production can block the event loop.
+See [batch writes](docs/batch-writes.md) for limits and recovery semantics.
+
 ## Prior art and alternatives
 
 `snowflake-sql-api` is not the first REST-based Snowflake client. If it does not

{snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.4}/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-0.1.3 → snowflake_sql_api-0.1.4}/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.3'
-__version_tuple__ = version_tuple = (0, 1, 3)
+__version__ = version = '0.1.4'
+__version_tuple__ = version_tuple = (0, 1, 4)
 
 __commit_id__ = commit_id = None

{snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.4}/snowflake_sql_api/aclient.py

@@ -17,6 +17,7 @@ from typing import (
     AsyncGenerator,
     Callable,
     Dict,
+    Iterable,
     List,
     Optional,
     Sequence,
@@ -27,6 +28,7 @@ from typing import (
 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
@@ -388,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-0.1.4/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)

{snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.4}/snowflake_sql_api/client.py

@@ -18,6 +18,7 @@ from typing import (
     Callable,
     Dict,
     Generator,
+    Iterable,
     List,
     Optional,
     Sequence,
@@ -28,6 +29,7 @@ from typing import (
 import httpx
 
 from .auth import KeypairAuthenticator, account_hostname
+from .batches import BatchWriteResult, _BatchExecution, _BatchPlanner
 from .bindings import to_bindings
 from .escaping import quote_identifier, quote_name
 from .exceptions import (
@@ -59,7 +61,16 @@ OnQuery = Callable[[str, "Params"], None]
 DEFAULT_POLL_TIMEOUT = 600.0
 
 
-def _rows_affected(body: Dict[str, Any]) -> int:
+def _affected_count(value: Any, strict: bool) -> int:
+    if strict and (isinstance(value, bool) or not isinstance(value, (int, str))):
+        raise ValueError("rows-affected count must be an integer")
+    count = int(value)
+    if strict and count < 0:
+        raise ValueError("rows-affected count must be nonnegative")
+    return count
+
+
+def _rows_affected(body: Dict[str, Any], *, strict: bool = False) -> int:
     """Extract the rows-affected count from a DML response."""
     stats = body.get("stats")
     if isinstance(stats, dict):
@@ -67,16 +78,19 @@ def _rows_affected(body: Dict[str, Any]) -> int:
         found = False
         for key in ("numRowsInserted", "numRowsUpdated", "numRowsDeleted"):
             if key in stats:
-                total += int(stats[key])
+                total += _affected_count(stats[key], strict)
                 found = True
         if found:
             return total
     data = body.get("data")
     if data and data[0]:
         try:
-            return int(data[0][0])
+            return _affected_count(data[0][0], strict)
         except (ValueError, TypeError):
-            pass
+            if strict:
+                raise
+    if strict:
+        raise ValueError("response has no rows-affected count")
     return 0
 
 
@@ -421,6 +435,52 @@ class SnowflakeClient:
         sql = f"INSERT INTO {quote_name(table)} ({cols_sql}) VALUES {values_sql}"
         return self.execute(sql, flat)
 
+    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 = 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 = 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 submit -----------------------------------------------------
 
     def submit(self, sql: str, params: Params = None) -> QueryHandle:

{snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.4}/snowflake_sql_api/exceptions.py

@@ -7,10 +7,15 @@ HTTP responses onto these types; see :mod:`snowflake_sql_api.transport`.
 
 from __future__ import annotations
 
-from typing import Optional
+from functools import partial
+from typing import TYPE_CHECKING, Any, Optional
+
+if TYPE_CHECKING:
+    from .batches import BatchOutcome, BatchPhase, BatchWriteResult, FailedBatch
 
 __all__ = [
     "SnowflakeError",
+    "BatchWriteError",
     "SnowflakeConfigError",
     "SnowflakeAuthError",
     "SnowflakeRequestError",
@@ -23,7 +28,54 @@ __all__ = [
 
 
 class SnowflakeError(Exception):
-    """Base class for every error raised by this library."""
+    """Base class for every error raised by this library.
+
+    ``submission_attempts`` counts actual parent POST attempts, or is None for
+    errors outside statement submission (including GET requests).
+    """
+
+    def __init__(
+        self, *args: object, submission_attempts: Optional[int] = None
+    ) -> None:
+        super().__init__(*args)
+        self.submission_attempts = submission_attempts
+
+
+class BatchWriteError(SnowflakeError):
+    """A batch import stopped; inspect confirmed progress and reconcile unknown writes."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        completed: BatchWriteResult,
+        next_row_index: int,
+        consumed_rows: int,
+        failed_batch: Optional[FailedBatch],
+        phase: BatchPhase,
+        outcome: BatchOutcome,
+    ) -> None:
+        super().__init__(message)
+        self.completed = completed
+        self.next_row_index = next_row_index
+        self.consumed_rows = consumed_rows
+        self.failed_batch = failed_batch
+        self.phase = phase
+        self.outcome = outcome
+
+    def __reduce__(self) -> tuple[Any, ...]:
+        fields = {
+            key: getattr(self, key)
+            for key in (
+                "completed",
+                "next_row_index",
+                "consumed_rows",
+                "failed_batch",
+                "phase",
+                "outcome",
+            )
+        }
+        return partial(type(self), **fields), self.args, self.__dict__
 
 
 class SnowflakeConfigError(SnowflakeError):
@@ -35,7 +87,26 @@ class SnowflakeConfigError(SnowflakeError):
 
 
 class SnowflakeAuthError(SnowflakeError):
-    """Keypair JWT generation or Snowflake authentication failed."""
+    """Snowflake rejected authentication (HTTP 401).
+
+    Carries the server envelope. Local JWT signing errors propagate with their
+    original types instead of being converted to this exception.
+    """
+
+    def __init__(
+        self,
+        *args: object,
+        status_code: Optional[int] = None,
+        code: Optional[str] = None,
+        sql_state: Optional[str] = None,
+        statement_handle: Optional[str] = None,
+        submission_attempts: Optional[int] = None,
+    ) -> None:
+        super().__init__(*args, submission_attempts=submission_attempts)
+        self.status_code = status_code
+        self.code = code
+        self.sql_state = sql_state
+        self.statement_handle = statement_handle
 
 
 class SnowflakeRequestError(SnowflakeError):
@@ -54,8 +125,9 @@ class SnowflakeRequestError(SnowflakeError):
         sql_state: Optional[str] = None,
         request_id: Optional[str] = None,
         statement_handle: Optional[str] = None,
+        submission_attempts: Optional[int] = None,
     ) -> None:
-        super().__init__(message)
+        super().__init__(message, submission_attempts=submission_attempts)
         self.status_code = status_code
         self.code = code
         self.sql_state = sql_state