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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (48) hide show
  1. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/PKG-INFO +48 -1
  2. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/README.md +47 -0
  3. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/__init__.py +6 -0
  4. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/_version.py +2 -2
  5. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/aclient.py +106 -9
  6. snowflake_sql_api-0.1.5/snowflake_sql_api/batches.py +298 -0
  7. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/client.py +133 -12
  8. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/exceptions.py +78 -4
  9. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/retry.py +5 -5
  10. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/transport.py +132 -53
  11. snowflake_sql_api-0.1.5/tests/smoke/test_submissions.py +300 -0
  12. snowflake_sql_api-0.1.5/tests/test_batch_writes.py +1023 -0
  13. snowflake_sql_api-0.1.5/tests/test_submissions.py +519 -0
  14. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/.gitignore +0 -0
  15. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/CONTRIBUTING.md +0 -0
  16. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/LICENSE +0 -0
  17. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/pyproject.toml +0 -0
  18. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/auth.py +0 -0
  19. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/bindings.py +0 -0
  20. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/cli.py +0 -0
  21. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/escaping.py +0 -0
  22. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/exports.py +0 -0
  23. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/pagination.py +0 -0
  24. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/py.typed +0 -0
  25. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/row_mapping.py +0 -0
  26. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/testing.py +0 -0
  27. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/types.py +0 -0
  28. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/__init__.py +0 -0
  29. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/conftest.py +0 -0
  30. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/smoke/__init__.py +0 -0
  31. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/smoke/test_smoke.py +0 -0
  32. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/streaming_support.py +0 -0
  33. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/support.py +0 -0
  34. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_aclient.py +0 -0
  35. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_auth.py +0 -0
  36. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_bindings.py +0 -0
  37. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_cli.py +0 -0
  38. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_client.py +0 -0
  39. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_escaping.py +0 -0
  40. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_exports.py +0 -0
  41. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_import_footprint.py +0 -0
  42. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_package.py +0 -0
  43. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_regression.py +0 -0
  44. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_retry.py +0 -0
  45. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_streaming.py +0 -0
  46. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_testing.py +0 -0
  47. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_transport.py +0 -0
  48. {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_types.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: snowflake-sql-api
3
- Version: 0.1.3
3
+ Version: 0.1.5
4
4
  Summary: Lightweight, pure-Python client for Snowflake's SQL API v2
5
5
  Project-URL: Homepage, https://github.com/hampsterx/snowflake-sql-api
6
6
  Project-URL: Repository, https://github.com/hampsterx/snowflake-sql-api
@@ -95,6 +95,7 @@ stays small and cold-starts quickly. Core dependencies: `httpx`, `PyJWT`,
95
95
  - Multi-partition result handling for large result sets
96
96
  - Lazy row streams and bounded row batches, with CSV and JSONL exports
97
97
  - Query helpers, DML, batch insert, async statement submission
98
+ - Caller-owned submission UUIDs and explicit replay across worker restarts
98
99
  - A small CLI for ad-hoc queries
99
100
  - Optional pandas (`[pandas]`) and typed-row (`[pydantic]`) helpers, kept out of
100
101
  the default install
@@ -155,6 +156,52 @@ with the largest decoded partition plus the application batch. See
155
156
  [streaming results and exports](docs/streaming.md) for row streams, async usage,
156
157
  serialization, and partial-output behavior.
157
158
 
159
+ Insert a record stream in bounded requests:
160
+
161
+ ```python
162
+ with SnowflakeClient.from_env() as client:
163
+ progress = client.insert_batches(
164
+ "events", ["id", "label"], ((i, f"event {i}") for i in range(2500)),
165
+ max_rows=1000, max_bytes=1_000_000,
166
+ )
167
+ print(progress.input_rows, progress.rows_affected)
168
+ ```
169
+
170
+ `insert_many()` sends one statement; `insert_batches()` sends sequential statements
171
+ and can partially complete. Its byte budget covers the exact JSON body. On failure,
172
+ `BatchWriteError` reports confirmed progress, consumed input and the failed request's
173
+ identity. Reconstruct replayable input and reconcile unknown outcomes before retrying;
174
+ do not resume with the same advanced iterator. The async equivalent is awaited and
175
+ accepts the same synchronous iterable, whose production can block the event loop.
176
+ See [batch writes](docs/batch-writes.md) for limits and recovery semantics.
177
+
178
+ ## Replayable submissions
179
+
180
+ For a queued job, allocate a UUID and persist it with the SQL, typed inputs and
181
+ explicit execution context **before** calling `submit()`. Pass that saved ID:
182
+
183
+ ```python
184
+ handle = client.submit(job["sql"], job["params"], request_id=job["request_id"])
185
+ # Persist handle.statement_handle when it reaches the worker.
186
+
187
+ # After a restart with no saved statement handle, rebuild identical inputs:
188
+ handle = client.resubmit(job["sql"], job["params"], request_id=job["request_id"])
189
+ rows = handle.result() # Only then mark the durable job complete.
190
+ ```
191
+
192
+ `resubmit()` sends the same ID with `retry=true` on its first HTTP attempt.
193
+ A saved statement handle can instead be attached with `QueryHandle(client, saved_handle)`
194
+ and polled without another submission. Both clients expose `handle.request_id`, and
195
+ submission errors in the `SnowflakeError` family carry the allocated ID. Async
196
+ submission, resubmission and result methods are awaited.
197
+
198
+ The application owns durable storage, identical payload reconstruction, UUID
199
+ uniqueness and reconciliation. Retained Snowflake history governs recovery; there
200
+ is no unlimited replay window or exactly-once processing guarantee. Live restart
201
+ verification requires the opt-in disposable-resource probes. See
202
+ [replayable submissions](docs/submissions.md) for the complete persistence recipe,
203
+ limits and verification status.
204
+
158
205
  ## Prior art and alternatives
159
206
 
160
207
  `snowflake-sql-api` is not the first REST-based Snowflake client. If it does not
@@ -31,6 +31,7 @@ stays small and cold-starts quickly. Core dependencies: `httpx`, `PyJWT`,
31
31
  - Multi-partition result handling for large result sets
32
32
  - Lazy row streams and bounded row batches, with CSV and JSONL exports
33
33
  - Query helpers, DML, batch insert, async statement submission
34
+ - Caller-owned submission UUIDs and explicit replay across worker restarts
34
35
  - A small CLI for ad-hoc queries
35
36
  - Optional pandas (`[pandas]`) and typed-row (`[pydantic]`) helpers, kept out of
36
37
  the default install
@@ -91,6 +92,52 @@ with the largest decoded partition plus the application batch. See
91
92
  [streaming results and exports](docs/streaming.md) for row streams, async usage,
92
93
  serialization, and partial-output behavior.
93
94
 
95
+ Insert a record stream in bounded requests:
96
+
97
+ ```python
98
+ with SnowflakeClient.from_env() as client:
99
+ progress = client.insert_batches(
100
+ "events", ["id", "label"], ((i, f"event {i}") for i in range(2500)),
101
+ max_rows=1000, max_bytes=1_000_000,
102
+ )
103
+ print(progress.input_rows, progress.rows_affected)
104
+ ```
105
+
106
+ `insert_many()` sends one statement; `insert_batches()` sends sequential statements
107
+ and can partially complete. Its byte budget covers the exact JSON body. On failure,
108
+ `BatchWriteError` reports confirmed progress, consumed input and the failed request's
109
+ identity. Reconstruct replayable input and reconcile unknown outcomes before retrying;
110
+ do not resume with the same advanced iterator. The async equivalent is awaited and
111
+ accepts the same synchronous iterable, whose production can block the event loop.
112
+ See [batch writes](docs/batch-writes.md) for limits and recovery semantics.
113
+
114
+ ## Replayable submissions
115
+
116
+ For a queued job, allocate a UUID and persist it with the SQL, typed inputs and
117
+ explicit execution context **before** calling `submit()`. Pass that saved ID:
118
+
119
+ ```python
120
+ handle = client.submit(job["sql"], job["params"], request_id=job["request_id"])
121
+ # Persist handle.statement_handle when it reaches the worker.
122
+
123
+ # After a restart with no saved statement handle, rebuild identical inputs:
124
+ handle = client.resubmit(job["sql"], job["params"], request_id=job["request_id"])
125
+ rows = handle.result() # Only then mark the durable job complete.
126
+ ```
127
+
128
+ `resubmit()` sends the same ID with `retry=true` on its first HTTP attempt.
129
+ A saved statement handle can instead be attached with `QueryHandle(client, saved_handle)`
130
+ and polled without another submission. Both clients expose `handle.request_id`, and
131
+ submission errors in the `SnowflakeError` family carry the allocated ID. Async
132
+ submission, resubmission and result methods are awaited.
133
+
134
+ The application owns durable storage, identical payload reconstruction, UUID
135
+ uniqueness and reconciliation. Retained Snowflake history governs recovery; there
136
+ is no unlimited replay window or exactly-once processing guarantee. Live restart
137
+ verification requires the opt-in disposable-resource probes. See
138
+ [replayable submissions](docs/submissions.md) for the complete persistence recipe,
139
+ limits and verification status.
140
+
94
141
  ## Prior art and alternatives
95
142
 
96
143
  `snowflake-sql-api` is not the first REST-based Snowflake client. If it does not
@@ -27,8 +27,10 @@ except ImportError: # pragma: no cover - plain source checkout, never built
27
27
  __version__ = "0.0.0.dev0"
28
28
 
29
29
  from .aclient import AsyncQueryHandle, AsyncSnowflakeClient
30
+ from .batches import BatchSummary, BatchWriteResult, FailedBatch
30
31
  from .client import QueryHandle, SnowflakeClient
31
32
  from .exceptions import (
33
+ BatchWriteError,
32
34
  ResultNotReady,
33
35
  SnowflakeAuthError,
34
36
  SnowflakeConfigError,
@@ -48,6 +50,10 @@ __all__ = [
48
50
  "QueryHandle",
49
51
  "AsyncQueryHandle",
50
52
  "RetryPolicy",
53
+ "BatchSummary",
54
+ "BatchWriteResult",
55
+ "FailedBatch",
56
+ "BatchWriteError",
51
57
  "SnowflakeError",
52
58
  "SnowflakeConfigError",
53
59
  "SnowflakeAuthError",
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
18
18
  commit_id: str | None
19
19
  __commit_id__: str | None
20
20
 
21
- __version__ = version = '0.1.3'
22
- __version_tuple__ = version_tuple = (0, 1, 3)
21
+ __version__ = version = '0.1.5'
22
+ __version_tuple__ = version_tuple = (0, 1, 5)
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -17,6 +17,7 @@ from typing import (
17
17
  AsyncGenerator,
18
18
  Callable,
19
19
  Dict,
20
+ Iterable,
20
21
  List,
21
22
  Optional,
22
23
  Sequence,
@@ -27,12 +28,14 @@ from typing import (
27
28
  import httpx
28
29
 
29
30
  from .auth import KeypairAuthenticator, account_hostname
31
+ from .batches import BatchWriteResult, _BatchExecution, _BatchPlanner
30
32
  from .bindings import to_bindings
31
- from .client import DEFAULT_POLL_TIMEOUT, _columns_from, _rows_affected
33
+ from .client import DEFAULT_POLL_TIMEOUT, _columns_from, _rows_affected, _submission_id
32
34
  from .escaping import quote_identifier, quote_name
33
35
  from .exceptions import (
34
36
  ResultNotReady,
35
37
  SnowflakeConfigError,
38
+ SnowflakeError,
36
39
  SnowflakeRequestError,
37
40
  SnowflakeTimeoutError,
38
41
  )
@@ -59,9 +62,16 @@ OnQuery = Callable[[str, "Params"], None]
59
62
  class AsyncQueryHandle:
60
63
  """Async handle to a submitted statement (see :class:`~snowflake_sql_api.client.QueryHandle`)."""
61
64
 
62
- def __init__(self, client: AsyncSnowflakeClient, statement_handle: str) -> None:
65
+ def __init__(
66
+ self,
67
+ client: AsyncSnowflakeClient,
68
+ statement_handle: str,
69
+ *,
70
+ request_id: Optional[str] = None,
71
+ ) -> None:
63
72
  self._client = client
64
73
  self.statement_handle = statement_handle
74
+ self.request_id = request_id
65
75
 
66
76
  async def result(
67
77
  self, *, poll: bool = True, timeout: Optional[float] = None
@@ -388,13 +398,100 @@ class AsyncSnowflakeClient:
388
398
 
389
399
  # -- async submit -----------------------------------------------------
390
400
 
391
- async def submit(self, sql: str, params: Params = None) -> AsyncQueryHandle:
392
- """Submit a long-running statement and return an :class:`AsyncQueryHandle`."""
393
- response = await self._run(sql, params, async_exec=True)
394
- handle = response.statement_handle
395
- if handle is None:
396
- raise SnowflakeRequestError("async submit returned no statement handle")
397
- return AsyncQueryHandle(self, handle)
401
+ async def insert_batches(
402
+ self,
403
+ table: str,
404
+ columns: Sequence[str],
405
+ rows: Iterable[Sequence[Any]],
406
+ *,
407
+ max_rows: int = 1000,
408
+ max_bytes: int = 1_000_000,
409
+ ) -> BatchWriteResult:
410
+ """Insert a single-pass source sequentially within row and wire-byte budgets.
411
+
412
+ Independent statements can commit before a later failure. BatchWriteError
413
+ reports confirmed progress and uncertain outcomes; see docs/batch-writes.md.
414
+ Input is an ordinary iterable, consumed on the calling thread.
415
+ """
416
+ planner = _BatchPlanner(
417
+ table, columns, rows, max_rows, max_bytes, self._build_payload
418
+ )
419
+ while (batch := planner.next_batch()) is not None:
420
+ try:
421
+ self._notify(batch.sql, batch.params)
422
+ except Exception as exc:
423
+ raise planner.error(exc, "validation") from exc
424
+ execution = _BatchExecution(planner, batch)
425
+ try:
426
+ response = await self._transport.submit(
427
+ {}, request_id=execution.failed.request_id, content=batch.content
428
+ )
429
+ execution.received(
430
+ response.statement_handle,
431
+ response.submission_attempts,
432
+ response.is_running,
433
+ )
434
+ if response.is_running:
435
+ handle = response.statement_handle
436
+ if handle is None:
437
+ raise SnowflakeRequestError("running statement has no handle")
438
+ response = await self._poll(handle, self._statement_timeout)
439
+ execution.succeeded = True
440
+ affected = _rows_affected(response.body, strict=True)
441
+ except Exception as exc:
442
+ raise execution.error(exc) from exc
443
+ planner.confirm(batch, affected)
444
+ del batch
445
+ return planner.result()
446
+
447
+ async def submit(
448
+ self,
449
+ sql: str,
450
+ params: Params = None,
451
+ *,
452
+ request_id: str | uuid.UUID | None = None,
453
+ ) -> AsyncQueryHandle:
454
+ """Submit once with a caller-owned UUID, or generate a fresh UUID4.
455
+
456
+ Persist the ID and identical inputs before calling for restart recovery.
457
+ Use :meth:`resubmit` once a previous attempt might have reached Snowflake.
458
+ """
459
+ return await self._submit_handle(
460
+ sql, params, request_id=request_id, retry=False
461
+ )
462
+
463
+ async def resubmit(
464
+ self, sql: str, params: Params = None, *, request_id: str | uuid.UUID
465
+ ) -> AsyncQueryHandle:
466
+ """Replay identical saved inputs with the same UUID and ``retry=true``.
467
+
468
+ The caller owns durable storage, payload consistency and reconciliation
469
+ when Snowflake's retained history cannot resolve an uncertain outcome.
470
+ """
471
+ return await self._submit_handle(sql, params, request_id=request_id, retry=True)
472
+
473
+ async def _submit_handle(
474
+ self,
475
+ sql: str,
476
+ params: Params,
477
+ *,
478
+ request_id: str | uuid.UUID | None,
479
+ retry: bool,
480
+ ) -> AsyncQueryHandle:
481
+ canonical_id = _submission_id(request_id, retry=retry)
482
+ self._notify(sql, params)
483
+ try:
484
+ payload = self._build_payload(sql, params)
485
+ response = await self._transport.submit(
486
+ payload, request_id=canonical_id, async_exec=True, retry=retry
487
+ )
488
+ handle = response.statement_handle
489
+ if handle is None:
490
+ raise SnowflakeRequestError("async submit returned no statement handle")
491
+ return AsyncQueryHandle(self, handle, request_id=canonical_id)
492
+ except SnowflakeError as exc:
493
+ exc.request_id = canonical_id
494
+ raise
398
495
 
399
496
  # -- lifecycle --------------------------------------------------------
400
497
 
@@ -0,0 +1,298 @@
1
+ """Bounded INSERT planning and immutable progress shared by both clients."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import uuid
7
+ from collections.abc import Iterable, Iterator, Sequence
8
+ from dataclasses import dataclass, replace
9
+ from typing import Any, Callable, Literal
10
+
11
+ from .bindings import to_bindings
12
+ from .escaping import quote_identifier, quote_name
13
+ from .exceptions import (
14
+ BatchWriteError,
15
+ SnowflakeAuthError,
16
+ SnowflakeConfigError,
17
+ SnowflakeProgrammingError,
18
+ SnowflakeRequestError,
19
+ )
20
+
21
+ __all__ = ["BatchSummary", "BatchWriteResult", "FailedBatch"]
22
+
23
+ BatchPhase = Literal["input", "validation", "submission", "execution"]
24
+ BatchOutcome = Literal["not_executed", "failed", "unknown", "succeeded"]
25
+
26
+
27
+ @dataclass(frozen=True)
28
+ class BatchSummary:
29
+ """One confirmed INSERT, with zero-based batch and source offsets."""
30
+
31
+ index: int
32
+ start_row_index: int
33
+ row_count: int
34
+ rows_affected: int
35
+
36
+
37
+ @dataclass(frozen=True)
38
+ class BatchWriteResult:
39
+ """Confirmed progress only; never retains input records or request bodies."""
40
+
41
+ batches: tuple[BatchSummary, ...] = ()
42
+
43
+ @property
44
+ def input_rows(self) -> int:
45
+ return sum(batch.row_count for batch in self.batches)
46
+
47
+ @property
48
+ def rows_affected(self) -> int:
49
+ return sum(batch.rows_affected for batch in self.batches)
50
+
51
+ @property
52
+ def completed_batches(self) -> int:
53
+ return len(self.batches)
54
+
55
+
56
+ @dataclass(frozen=True)
57
+ class FailedBatch:
58
+ """Prepared INSERT identity; an unknown outcome requires reconciliation."""
59
+
60
+ index: int
61
+ start_row_index: int
62
+ row_count: int
63
+ request_id: str
64
+ statement_handle: str | None = None
65
+ submission_attempts: int = 0
66
+
67
+
68
+ def _encode(value: Any) -> bytes:
69
+ return json.dumps(
70
+ value, ensure_ascii=False, separators=(",", ":"), allow_nan=False
71
+ ).encode("utf-8")
72
+
73
+
74
+ @dataclass
75
+ class _PreparedBatch:
76
+ sql: str
77
+ params: list[Any]
78
+ content: bytes
79
+ row_count: int
80
+
81
+
82
+ class _BatchPlanner:
83
+ def __init__(
84
+ self,
85
+ table: str,
86
+ columns: Sequence[str],
87
+ rows: Iterable[Sequence[Any]],
88
+ max_rows: int,
89
+ max_bytes: int,
90
+ build_payload: Callable[[str, Any], dict[str, Any]],
91
+ ) -> None:
92
+ if not columns:
93
+ raise SnowflakeConfigError("insert_batches requires at least one column")
94
+ for name, limit in (("max_rows", max_rows), ("max_bytes", max_bytes)):
95
+ if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0:
96
+ raise SnowflakeConfigError(f"{name} must be a positive integer")
97
+ if max_rows > 200_000:
98
+ raise SnowflakeConfigError("max_rows must be <= 200000")
99
+ cols = ", ".join(quote_identifier(c) for c in columns)
100
+ self.prefix = f"INSERT INTO {quote_name(table)} ({cols}) VALUES "
101
+ self.placeholder = "(" + ", ".join(["?"] * len(columns)) + ")"
102
+ self.width = len(columns)
103
+ self.max_rows, self.max_bytes = max_rows, max_bytes
104
+ self.source = rows
105
+ self.iterator: Iterator[Sequence[Any]] | None = None
106
+ self.build_payload = build_payload
107
+ self.context: bytes | None = None
108
+ self.single_overhead = 0
109
+ self.pending: tuple[list[Any], list[bytes]] | None = None
110
+ self.exhausted = False
111
+ self.consumed_rows = 0
112
+ self.confirmed_rows = 0
113
+ self.completed: list[BatchSummary] = []
114
+
115
+ def result(self) -> BatchWriteResult:
116
+ return BatchWriteResult(tuple(self.completed))
117
+
118
+ def error(
119
+ self,
120
+ cause: Exception,
121
+ phase: BatchPhase,
122
+ outcome: BatchOutcome = "not_executed",
123
+ failed_batch: FailedBatch | None = None,
124
+ ) -> BatchWriteError:
125
+ return BatchWriteError(
126
+ f"batch write stopped during {phase}: {cause}",
127
+ completed=self.result(),
128
+ next_row_index=self.confirmed_rows,
129
+ consumed_rows=self.consumed_rows,
130
+ failed_batch=failed_batch,
131
+ phase=phase,
132
+ outcome=outcome,
133
+ )
134
+
135
+ def _record(self) -> tuple[list[Any], list[bytes]] | None:
136
+ if self.exhausted:
137
+ return None
138
+ try:
139
+ if self.iterator is None:
140
+ self.iterator = iter(self.source)
141
+ except Exception as exc:
142
+ raise self.error(exc, "input") from exc
143
+ try:
144
+ row = next(self.iterator)
145
+ except StopIteration:
146
+ self.exhausted = True
147
+ return None
148
+ except Exception as exc:
149
+ raise self.error(exc, "input") from exc
150
+ self.consumed_rows += 1
151
+ try:
152
+ if not isinstance(row, Sequence) or isinstance(
153
+ row, (str, bytes, bytearray)
154
+ ):
155
+ raise SnowflakeConfigError("each row must be a positional sequence")
156
+ if len(row) != self.width:
157
+ raise SnowflakeConfigError(
158
+ f"each row must have {self.width} values, got {len(row)}"
159
+ )
160
+ values = list(row)
161
+ return values, [_encode(value) for value in to_bindings(values).values()]
162
+ except Exception as exc:
163
+ raise self.error(exc, "validation") from exc
164
+
165
+ @staticmethod
166
+ def _fragment(values: list[bytes], offset: int) -> bytes:
167
+ return b",".join(
168
+ _encode(str(offset + i)) + b":" + value for i, value in enumerate(values, 1)
169
+ )
170
+
171
+ def _body(self, sql: str, bindings: bytes) -> bytes:
172
+ assert self.context is not None
173
+ return (
174
+ b'{"statement":'
175
+ + _encode(sql)
176
+ + self.context
177
+ + b',"bindings":{'
178
+ + bindings
179
+ + b"}}"
180
+ )
181
+
182
+ def next_batch(self) -> _PreparedBatch | None:
183
+ # Configuration is validated before this point; input errors discard the
184
+ # unsent buffer. Only a validated singleton can trigger a byte-boundary flush.
185
+ params: list[Any] = []
186
+ fragments: list[bytes] = []
187
+ row_count = 0
188
+ size = 0
189
+ while row_count < self.max_rows:
190
+ record = self.pending
191
+ self.pending = None
192
+ if record is None:
193
+ record = self._record()
194
+ if record is None:
195
+ break
196
+ values, encoded = record
197
+ try:
198
+ if self.context is None:
199
+ context = self.build_payload("", None)
200
+ del context["statement"]
201
+ self.context = _encode(context)[1:-1]
202
+ if self.context:
203
+ self.context = b"," + self.context
204
+ self.single_overhead = len(
205
+ self._body(self.prefix + self.placeholder, b"")
206
+ )
207
+ first = self._fragment(encoded, 0)
208
+ single_size = self.single_overhead + len(first)
209
+ if single_size > self.max_bytes:
210
+ raise SnowflakeConfigError(
211
+ f"row {self.consumed_rows - 1} exceeds max_bytes={self.max_bytes}"
212
+ )
213
+ fragment = (
214
+ first if row_count == 0 else self._fragment(encoded, len(params))
215
+ )
216
+ candidate = (
217
+ single_size
218
+ if row_count == 0
219
+ else size + 2 + len(self.placeholder) + 1 + len(fragment)
220
+ )
221
+ if candidate > self.max_bytes:
222
+ self.pending = record
223
+ break
224
+ fragments.append(fragment)
225
+ params.extend(values)
226
+ row_count += 1
227
+ size = candidate
228
+ except Exception as exc:
229
+ raise self.error(exc, "validation") from exc
230
+ if not row_count:
231
+ return None
232
+ sql = self.prefix + ", ".join([self.placeholder] * row_count)
233
+ content = self._body(sql, b",".join(fragments))
234
+ if len(content) != size or len(content) > self.max_bytes:
235
+ budget_error = SnowflakeConfigError(
236
+ "batch body exceeds its measured byte budget"
237
+ )
238
+ raise self.error(budget_error, "validation") from budget_error
239
+ return _PreparedBatch(sql, params, content, row_count)
240
+
241
+ def confirm(self, batch: _PreparedBatch, rows_affected: int) -> None:
242
+ self.completed.append(
243
+ BatchSummary(
244
+ len(self.completed), self.confirmed_rows, batch.row_count, rows_affected
245
+ )
246
+ )
247
+ self.confirmed_rows += batch.row_count
248
+
249
+
250
+ class _BatchExecution:
251
+ """Parent POST context survives polling and conservative error classification."""
252
+
253
+ def __init__(self, planner: _BatchPlanner, batch: _PreparedBatch) -> None:
254
+ self.planner = planner
255
+ self.failed = FailedBatch(
256
+ len(planner.completed),
257
+ planner.confirmed_rows,
258
+ batch.row_count,
259
+ str(uuid.uuid4()),
260
+ )
261
+ self.phase: BatchPhase = "submission"
262
+ self.succeeded = False
263
+
264
+ def received(self, handle: str | None, attempts: int | None, running: bool) -> None:
265
+ self.failed = replace(
266
+ self.failed, statement_handle=handle, submission_attempts=attempts or 0
267
+ )
268
+ if not running or handle is not None:
269
+ self.phase = "execution"
270
+ self.succeeded = not running
271
+
272
+ def error(self, cause: Exception) -> BatchWriteError:
273
+ attempts = getattr(cause, "submission_attempts", None)
274
+ handle = getattr(cause, "statement_handle", None) or getattr(
275
+ cause.__cause__, "statement_handle", None
276
+ )
277
+ self.failed = replace(
278
+ self.failed,
279
+ submission_attempts=(
280
+ self.failed.submission_attempts if attempts is None else attempts
281
+ ),
282
+ statement_handle=handle or self.failed.statement_handle,
283
+ )
284
+ outcome: BatchOutcome = "succeeded" if self.succeeded else "unknown"
285
+ if not self.succeeded and self.failed.submission_attempts <= 1:
286
+ if isinstance(cause, SnowflakeProgrammingError):
287
+ self.phase, outcome = "execution", "failed"
288
+ elif self.phase == "submission" and (
289
+ attempts == 0
290
+ or isinstance(cause, SnowflakeAuthError)
291
+ and cause.status_code == 401
292
+ or isinstance(cause, SnowflakeRequestError)
293
+ and cause.status_code in (400, 401, 403)
294
+ ):
295
+ outcome = "not_executed"
296
+ if isinstance(cause, SnowflakeProgrammingError):
297
+ self.phase = "execution"
298
+ return self.planner.error(cause, self.phase, outcome, self.failed)