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.
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/PKG-INFO +48 -1
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/README.md +47 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/__init__.py +6 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/_version.py +2 -2
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/aclient.py +106 -9
- snowflake_sql_api-0.1.5/snowflake_sql_api/batches.py +298 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/client.py +133 -12
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/exceptions.py +78 -4
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/retry.py +5 -5
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/transport.py +132 -53
- snowflake_sql_api-0.1.5/tests/smoke/test_submissions.py +300 -0
- snowflake_sql_api-0.1.5/tests/test_batch_writes.py +1023 -0
- snowflake_sql_api-0.1.5/tests/test_submissions.py +519 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/.gitignore +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/CONTRIBUTING.md +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/LICENSE +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/pyproject.toml +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/auth.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/bindings.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/cli.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/escaping.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/exports.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/pagination.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/py.typed +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/row_mapping.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/testing.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/snowflake_sql_api/types.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/__init__.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/conftest.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/smoke/__init__.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/smoke/test_smoke.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/streaming_support.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/support.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_aclient.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_auth.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_bindings.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_cli.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_client.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_escaping.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_exports.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_import_footprint.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_package.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_regression.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_retry.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_streaming.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_testing.py +0 -0
- {snowflake_sql_api-0.1.3 → snowflake_sql_api-0.1.5}/tests/test_transport.py +0 -0
- {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
|
+
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.
|
|
22
|
-
__version_tuple__ = version_tuple = (0, 1,
|
|
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__(
|
|
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
|
|
392
|
-
|
|
393
|
-
|
|
394
|
-
|
|
395
|
-
|
|
396
|
-
|
|
397
|
-
|
|
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)
|