snowflake-sql-api 0.1.2__tar.gz → 0.1.4__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 (47) hide show
  1. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/CONTRIBUTING.md +1 -1
  2. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/PKG-INFO +52 -10
  3. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/README.md +47 -5
  4. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/pyproject.toml +17 -19
  5. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/__init__.py +6 -0
  6. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/_version.py +2 -2
  7. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/aclient.py +156 -6
  8. snowflake_sql_api-0.1.4/snowflake_sql_api/batches.py +298 -0
  9. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/client.py +160 -10
  10. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/exceptions.py +76 -4
  11. snowflake_sql_api-0.1.4/snowflake_sql_api/exports.py +71 -0
  12. snowflake_sql_api-0.1.4/snowflake_sql_api/pagination.py +131 -0
  13. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/transport.py +120 -48
  14. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/types.py +10 -6
  15. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/conftest.py +6 -0
  16. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/smoke/test_smoke.py +37 -0
  17. snowflake_sql_api-0.1.4/tests/streaming_support.py +73 -0
  18. snowflake_sql_api-0.1.4/tests/test_batch_writes.py +1023 -0
  19. snowflake_sql_api-0.1.4/tests/test_exports.py +221 -0
  20. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_import_footprint.py +1 -3
  21. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_regression.py +2 -2
  22. snowflake_sql_api-0.1.4/tests/test_streaming.py +279 -0
  23. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_testing.py +0 -2
  24. snowflake_sql_api-0.1.2/snowflake_sql_api/pagination.py +0 -53
  25. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/.gitignore +0 -0
  26. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/LICENSE +0 -0
  27. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/auth.py +0 -0
  28. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/bindings.py +0 -0
  29. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/cli.py +0 -0
  30. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/escaping.py +0 -0
  31. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/py.typed +0 -0
  32. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/retry.py +0 -0
  33. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/row_mapping.py +0 -0
  34. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/snowflake_sql_api/testing.py +0 -0
  35. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/__init__.py +0 -0
  36. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/smoke/__init__.py +0 -0
  37. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/support.py +0 -0
  38. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_aclient.py +0 -0
  39. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_auth.py +0 -0
  40. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_bindings.py +0 -0
  41. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_cli.py +0 -0
  42. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_client.py +0 -0
  43. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_escaping.py +0 -0
  44. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_package.py +0 -0
  45. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_retry.py +0 -0
  46. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_transport.py +0 -0
  47. {snowflake_sql_api-0.1.2 → snowflake_sql_api-0.1.4}/tests/test_types.py +0 -0
@@ -34,7 +34,7 @@ All four must pass (and `pre-commit run --all-files` is clean). Use `coverage ru
34
34
  (`snowflake_sql_api.testing`), and `coverage run` starts tracing before that
35
35
  plugin is imported so import-time lines are measured correctly. The coverage gate
36
36
  is enforced (`fail_under = 89` in `pyproject.toml`); `coverage report` exits
37
- non-zero below it. CI runs the same checks across Python 3.9-3.13.
37
+ non-zero below it. CI runs the same checks across Python 3.11-3.15.
38
38
 
39
39
  Checklist:
40
40
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: snowflake-sql-api
3
- Version: 0.1.2
3
+ Version: 0.1.4
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
@@ -34,20 +34,20 @@ Classifier: Intended Audience :: Developers
34
34
  Classifier: License :: OSI Approved :: MIT License
35
35
  Classifier: Operating System :: OS Independent
36
36
  Classifier: Programming Language :: Python :: 3
37
- Classifier: Programming Language :: Python :: 3.9
38
- Classifier: Programming Language :: Python :: 3.10
39
37
  Classifier: Programming Language :: Python :: 3.11
40
38
  Classifier: Programming Language :: Python :: 3.12
41
39
  Classifier: Programming Language :: Python :: 3.13
40
+ Classifier: Programming Language :: Python :: 3.14
41
+ Classifier: Programming Language :: Python :: 3.15
42
42
  Classifier: Programming Language :: Python :: Implementation :: CPython
43
43
  Classifier: Topic :: Database
44
44
  Classifier: Typing :: Typed
45
- Requires-Python: >=3.9
45
+ Requires-Python: >=3.11
46
46
  Requires-Dist: cryptography>=3.4
47
47
  Requires-Dist: httpx>=0.24
48
48
  Requires-Dist: pyjwt>=2.4
49
49
  Provides-Extra: dev
50
- Requires-Dist: black<25.10,>=24.0; extra == 'dev'
50
+ Requires-Dist: black<26.6,>=26.5; extra == 'dev'
51
51
  Requires-Dist: build>=1.0; extra == 'dev'
52
52
  Requires-Dist: coverage[toml]>=7.0; extra == 'dev'
53
53
  Requires-Dist: mypy>=1.8; extra == 'dev'
@@ -64,10 +64,16 @@ Description-Content-Type: text/markdown
64
64
 
65
65
  # snowflake-sql-api
66
66
 
67
+ [![CI](https://github.com/hampsterx/snowflake-sql-api/actions/workflows/ci.yml/badge.svg)](https://github.com/hampsterx/snowflake-sql-api/actions/workflows/ci.yml)
68
+ [![PyPI](https://img.shields.io/pypi/v/snowflake-sql-api)](https://pypi.org/project/snowflake-sql-api/)
69
+ [![Python versions](https://img.shields.io/pypi/pyversions/snowflake-sql-api)](https://pypi.org/project/snowflake-sql-api/)
70
+ [![License: MIT](https://img.shields.io/pypi/l/snowflake-sql-api)](LICENSE)
71
+ [![Typed](https://img.shields.io/badge/typing-typed-blue)](https://peps.python.org/pep-0561/)
72
+
67
73
  A lightweight, pure-Python client for [Snowflake's SQL API v2](https://docs.snowflake.com/en/developer-guide/sql-api/index)
68
74
  (the `POST /api/v2/statements` REST endpoint).
69
75
 
70
- > **Status: in active development, not yet released.** The API may change before `v0.1.0`.
76
+ > **Status: pre-1.0.** The API may still change between minor versions.
71
77
 
72
78
  ## Why
73
79
 
@@ -87,6 +93,7 @@ stays small and cold-starts quickly. Core dependencies: `httpx`, `PyJWT`,
87
93
  - Keypair (JWT) authentication, including encrypted private keys
88
94
  - Result type coercion (NUMBER/Decimal, dates/timestamps, VARIANT, and more)
89
95
  - Multi-partition result handling for large result sets
96
+ - Lazy row streams and bounded row batches, with CSV and JSONL exports
90
97
  - Query helpers, DML, batch insert, async statement submission
91
98
  - A small CLI for ad-hoc queries
92
99
  - Optional pandas (`[pandas]`) and typed-row (`[pydantic]`) helpers, kept out of
@@ -132,6 +139,41 @@ CLI:
132
139
  snowflake-sql-api query "SELECT current_version()"
133
140
  ```
134
141
 
142
+ Stream large results or export them without collecting every row:
143
+
144
+ ```python
145
+ with SnowflakeClient.from_env() as client:
146
+ for batch in client.query_batches("SELECT id, amount FROM sales", batch_size=500):
147
+ process(batch)
148
+
149
+ with open("sales.csv", "w", encoding="utf-8", newline="") as output:
150
+ count = client.export_csv("SELECT id, amount FROM sales", output)
151
+ ```
152
+
153
+ Streaming fetches completed results one server partition at a time. Memory scales
154
+ with the largest decoded partition plus the application batch. See
155
+ [streaming results and exports](docs/streaming.md) for row streams, async usage,
156
+ serialization, and partial-output behavior.
157
+
158
+ Insert a record stream in bounded requests:
159
+
160
+ ```python
161
+ with SnowflakeClient.from_env() as client:
162
+ progress = client.insert_batches(
163
+ "events", ["id", "label"], ((i, f"event {i}") for i in range(2500)),
164
+ max_rows=1000, max_bytes=1_000_000,
165
+ )
166
+ print(progress.input_rows, progress.rows_affected)
167
+ ```
168
+
169
+ `insert_many()` sends one statement; `insert_batches()` sends sequential statements
170
+ and can partially complete. Its byte budget covers the exact JSON body. On failure,
171
+ `BatchWriteError` reports confirmed progress, consumed input and the failed request's
172
+ identity. Reconstruct replayable input and reconcile unknown outcomes before retrying;
173
+ do not resume with the same advanced iterator. The async equivalent is awaited and
174
+ accepts the same synchronous iterable, whose production can block the event loop.
175
+ See [batch writes](docs/batch-writes.md) for limits and recovery semantics.
176
+
135
177
  ## Prior art and alternatives
136
178
 
137
179
  `snowflake-sql-api` is not the first REST-based Snowflake client. If it does not
@@ -146,10 +188,10 @@ fit your needs, one of these might:
146
188
  | [`rb_snowflake_client`](https://github.com/rinsed-org/rb-snowflake-client) | Ruby | Mature pure-Ruby SQL API client (streaming, connection pooling). |
147
189
  | [`rsql`](https://github.com/theseus-rs/rsql) | Rust | Multi-database SQL CLI with a Snowflake SQL API driver. |
148
190
 
149
- **How `snowflake-sql-api` differs:** pure Python over `httpx` with no dependency on
150
- the official connector (small install, fast cold start), paired sync and async
151
- clients, and keypair auth, type coercion, partition handling, and a CLI as tested
152
- first-class behavior.
191
+ **How `snowflake-sql-api` differs:** pure Python over `httpx`, with no
192
+ dependency on the official connector (small install, fast cold start), paired
193
+ sync and async clients, keypair auth, type coercion, partition handling, and a
194
+ CLI as tested first-class behavior.
153
195
 
154
196
  Note: the similarly named `snowflake-sql-api-async` and `snowflake-sql-api-client`
155
197
  are separate, unaffiliated projects.
@@ -1,9 +1,15 @@
1
1
  # snowflake-sql-api
2
2
 
3
+ [![CI](https://github.com/hampsterx/snowflake-sql-api/actions/workflows/ci.yml/badge.svg)](https://github.com/hampsterx/snowflake-sql-api/actions/workflows/ci.yml)
4
+ [![PyPI](https://img.shields.io/pypi/v/snowflake-sql-api)](https://pypi.org/project/snowflake-sql-api/)
5
+ [![Python versions](https://img.shields.io/pypi/pyversions/snowflake-sql-api)](https://pypi.org/project/snowflake-sql-api/)
6
+ [![License: MIT](https://img.shields.io/pypi/l/snowflake-sql-api)](LICENSE)
7
+ [![Typed](https://img.shields.io/badge/typing-typed-blue)](https://peps.python.org/pep-0561/)
8
+
3
9
  A lightweight, pure-Python client for [Snowflake's SQL API v2](https://docs.snowflake.com/en/developer-guide/sql-api/index)
4
10
  (the `POST /api/v2/statements` REST endpoint).
5
11
 
6
- > **Status: in active development, not yet released.** The API may change before `v0.1.0`.
12
+ > **Status: pre-1.0.** The API may still change between minor versions.
7
13
 
8
14
  ## Why
9
15
 
@@ -23,6 +29,7 @@ stays small and cold-starts quickly. Core dependencies: `httpx`, `PyJWT`,
23
29
  - Keypair (JWT) authentication, including encrypted private keys
24
30
  - Result type coercion (NUMBER/Decimal, dates/timestamps, VARIANT, and more)
25
31
  - Multi-partition result handling for large result sets
32
+ - Lazy row streams and bounded row batches, with CSV and JSONL exports
26
33
  - Query helpers, DML, batch insert, async statement submission
27
34
  - A small CLI for ad-hoc queries
28
35
  - Optional pandas (`[pandas]`) and typed-row (`[pydantic]`) helpers, kept out of
@@ -68,6 +75,41 @@ CLI:
68
75
  snowflake-sql-api query "SELECT current_version()"
69
76
  ```
70
77
 
78
+ Stream large results or export them without collecting every row:
79
+
80
+ ```python
81
+ with SnowflakeClient.from_env() as client:
82
+ for batch in client.query_batches("SELECT id, amount FROM sales", batch_size=500):
83
+ process(batch)
84
+
85
+ with open("sales.csv", "w", encoding="utf-8", newline="") as output:
86
+ count = client.export_csv("SELECT id, amount FROM sales", output)
87
+ ```
88
+
89
+ Streaming fetches completed results one server partition at a time. Memory scales
90
+ with the largest decoded partition plus the application batch. See
91
+ [streaming results and exports](docs/streaming.md) for row streams, async usage,
92
+ serialization, and partial-output behavior.
93
+
94
+ Insert a record stream in bounded requests:
95
+
96
+ ```python
97
+ with SnowflakeClient.from_env() as client:
98
+ progress = client.insert_batches(
99
+ "events", ["id", "label"], ((i, f"event {i}") for i in range(2500)),
100
+ max_rows=1000, max_bytes=1_000_000,
101
+ )
102
+ print(progress.input_rows, progress.rows_affected)
103
+ ```
104
+
105
+ `insert_many()` sends one statement; `insert_batches()` sends sequential statements
106
+ and can partially complete. Its byte budget covers the exact JSON body. On failure,
107
+ `BatchWriteError` reports confirmed progress, consumed input and the failed request's
108
+ identity. Reconstruct replayable input and reconcile unknown outcomes before retrying;
109
+ do not resume with the same advanced iterator. The async equivalent is awaited and
110
+ accepts the same synchronous iterable, whose production can block the event loop.
111
+ See [batch writes](docs/batch-writes.md) for limits and recovery semantics.
112
+
71
113
  ## Prior art and alternatives
72
114
 
73
115
  `snowflake-sql-api` is not the first REST-based Snowflake client. If it does not
@@ -82,10 +124,10 @@ fit your needs, one of these might:
82
124
  | [`rb_snowflake_client`](https://github.com/rinsed-org/rb-snowflake-client) | Ruby | Mature pure-Ruby SQL API client (streaming, connection pooling). |
83
125
  | [`rsql`](https://github.com/theseus-rs/rsql) | Rust | Multi-database SQL CLI with a Snowflake SQL API driver. |
84
126
 
85
- **How `snowflake-sql-api` differs:** pure Python over `httpx` with no dependency on
86
- the official connector (small install, fast cold start), paired sync and async
87
- clients, and keypair auth, type coercion, partition handling, and a CLI as tested
88
- first-class behavior.
127
+ **How `snowflake-sql-api` differs:** pure Python over `httpx`, with no
128
+ dependency on the official connector (small install, fast cold start), paired
129
+ sync and async clients, keypair auth, type coercion, partition handling, and a
130
+ CLI as tested first-class behavior.
89
131
 
90
132
  Note: the similarly named `snowflake-sql-api-async` and `snowflake-sql-api-client`
91
133
  are separate, unaffiliated projects.
@@ -7,7 +7,7 @@ name = "snowflake-sql-api"
7
7
  dynamic = ["version"]
8
8
  description = "Lightweight, pure-Python client for Snowflake's SQL API v2"
9
9
  readme = "README.md"
10
- requires-python = ">=3.9"
10
+ requires-python = ">=3.11"
11
11
  license = { file = "LICENSE" }
12
12
  authors = [{ name = "hampsterx" }]
13
13
  keywords = ["snowflake", "sql", "rest", "api", "httpx", "serverless", "lambda"]
@@ -17,11 +17,11 @@ classifiers = [
17
17
  "License :: OSI Approved :: MIT License",
18
18
  "Operating System :: OS Independent",
19
19
  "Programming Language :: Python :: 3",
20
- "Programming Language :: Python :: 3.9",
21
- "Programming Language :: Python :: 3.10",
22
20
  "Programming Language :: Python :: 3.11",
23
21
  "Programming Language :: Python :: 3.12",
24
22
  "Programming Language :: Python :: 3.13",
23
+ "Programming Language :: Python :: 3.14",
24
+ "Programming Language :: Python :: 3.15",
25
25
  "Programming Language :: Python :: Implementation :: CPython",
26
26
  "Topic :: Database",
27
27
  "Typing :: Typed",
@@ -41,12 +41,10 @@ dev = [
41
41
  "pytest-asyncio>=0.21",
42
42
  "respx>=0.20",
43
43
  "ruff>=0.4",
44
- # Capped below 25.10 so CI formats with the same black the pre-commit hook
45
- # pins (rev 25.9.0 in .pre-commit-config.yaml). Move the two together by
46
- # hand; Dependabot ignores black (.github/dependabot.yml). Keep the cap at
47
- # or below 25.12 while Python 3.9 is supported: 25.12.0+ require 3.10, so
48
- # a higher cap resolves a different black per interpreter.
49
- "black>=24.0,<25.10",
44
+ # Held to 26.5.x so CI formats with the same black the pre-commit hook
45
+ # pins (rev 26.5.1 in .pre-commit-config.yaml). Move the two together by
46
+ # hand; Dependabot ignores black (.github/dependabot.yml).
47
+ "black>=26.5,<26.6",
50
48
  "mypy>=1.8",
51
49
  "build>=1.0",
52
50
  "pre-commit>=3.5",
@@ -89,13 +87,13 @@ include = [
89
87
 
90
88
  [tool.ruff]
91
89
  line-length = 88
92
- target-version = "py39"
90
+ target-version = "py311"
93
91
 
94
92
  [tool.ruff.lint]
95
93
  select = ["E", "F", "W", "I", "B", "C4", "SIM"]
96
- # E501 (line length) is owned by black. `UP` (pyupgrade) is omitted: it would
97
- # rewrite runtime type aliases like `Optional[Sequence[Any]]` to `X | None`,
98
- # which raises at import time on Python 3.9 (PEP 604 on typing generics is 3.10+).
94
+ # E501 (line length) is owned by black. `UP` (pyupgrade) is not enabled yet: the
95
+ # code predates the 3.11 floor and still spells `Optional`/`List`, so turning it
96
+ # on is a ~250-site annotation rewrite of its own.
99
97
  ignore = ["E501"]
100
98
 
101
99
  [tool.ruff.lint.per-file-ignores]
@@ -103,14 +101,14 @@ ignore = ["E501"]
103
101
 
104
102
  [tool.black]
105
103
  line-length = 88
106
- # Pin to py39 grammar: our code uses no newer syntax, and a single floor target
107
- # keeps black's AST safety check parseable by the interpreter running CI.
108
- target-version = ["py39"]
104
+ # Pin to the floor's grammar: a single floor target keeps black's AST safety
105
+ # check parseable by the interpreter running CI.
106
+ target-version = ["py311"]
109
107
 
110
108
  [tool.mypy]
111
- # Lowest target this mypy accepts; true 3.9 runtime compat is gated by the
112
- # pytest matrix (3.9-3.13), which imports every module under each interpreter.
113
- python_version = "3.10"
109
+ # The runtime floor; the pytest matrix (3.11-3.15) imports every module under
110
+ # each supported interpreter.
111
+ python_version = "3.11"
114
112
  files = ["snowflake_sql_api"]
115
113
  disallow_untyped_defs = true
116
114
  disallow_incomplete_defs = true
@@ -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.2'
22
- __version_tuple__ = version_tuple = (0, 1, 2)
21
+ __version__ = version = '0.1.4'
22
+ __version_tuple__ = version_tuple = (0, 1, 4)
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -12,11 +12,23 @@ import os
12
12
  import time
13
13
  import uuid
14
14
  from pathlib import Path
15
- from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
15
+ from typing import (
16
+ Any,
17
+ AsyncGenerator,
18
+ Callable,
19
+ Dict,
20
+ Iterable,
21
+ List,
22
+ Optional,
23
+ Sequence,
24
+ TextIO,
25
+ Tuple,
26
+ )
16
27
 
17
28
  import httpx
18
29
 
19
30
  from .auth import KeypairAuthenticator, account_hostname
31
+ from .batches import BatchWriteResult, _BatchExecution, _BatchPlanner
20
32
  from .bindings import to_bindings
21
33
  from .client import DEFAULT_POLL_TIMEOUT, _columns_from, _rows_affected
22
34
  from .escaping import quote_identifier, quote_name
@@ -26,7 +38,13 @@ from .exceptions import (
26
38
  SnowflakeRequestError,
27
39
  SnowflakeTimeoutError,
28
40
  )
29
- from .pagination import fetch_all_partitions_async
41
+ from .exports import (
42
+ CSVWriter,
43
+ validate_batch_size,
44
+ validate_export_columns,
45
+ write_jsonl,
46
+ )
47
+ from .pagination import fetch_all_partitions_async, iter_result_rows_async
30
48
  from .retry import DEFAULT_RETRY_POLICY, RetryPolicy
31
49
  from .transport import AsyncTransport, StatementResponse
32
50
  from .types import ColumnMeta, Row, coerce_rows
@@ -201,13 +219,13 @@ class AsyncSnowflakeClient:
201
219
  )
202
220
  await asyncio.sleep(self._poll_interval)
203
221
 
204
- async def _collect(
222
+ async def _prepare_result(
205
223
  self,
206
224
  response: StatementResponse,
207
225
  *,
208
226
  timeout: Optional[float],
209
227
  poll: bool = True,
210
- ) -> Tuple[List[ColumnMeta], List[List[Any]]]:
228
+ ) -> Tuple[List[ColumnMeta], str, Dict[str, Any]]:
211
229
  if response.is_running:
212
230
  if not poll:
213
231
  raise ResultNotReady(statement_handle=response.statement_handle)
@@ -220,8 +238,27 @@ class AsyncSnowflakeClient:
220
238
  handle = response.statement_handle
221
239
  if handle is None:
222
240
  raise SnowflakeRequestError("completed statement has no handle")
223
- raw_rows = await fetch_all_partitions_async(self._transport, handle, body)
224
- return columns, raw_rows
241
+ return columns, handle, body
242
+
243
+ async def _collect(
244
+ self,
245
+ response: StatementResponse,
246
+ *,
247
+ timeout: Optional[float],
248
+ poll: bool = True,
249
+ ) -> Tuple[List[ColumnMeta], List[List[Any]]]:
250
+ columns, handle, body = await self._prepare_result(
251
+ response, timeout=timeout, poll=poll
252
+ )
253
+ return columns, await fetch_all_partitions_async(self._transport, handle, body)
254
+
255
+ async def _stream_result(
256
+ self, sql: str, params: Params
257
+ ) -> Tuple[List[ColumnMeta], AsyncGenerator[Row, None]]:
258
+ columns, handle, body = await self._prepare_result(
259
+ await self._run(sql, params), timeout=self._statement_timeout
260
+ )
261
+ return columns, iter_result_rows_async(self._transport, handle, body, columns)
225
262
 
226
263
  # -- query surface ----------------------------------------------------
227
264
 
@@ -233,6 +270,73 @@ class AsyncSnowflakeClient:
233
270
  )
234
271
  return coerce_rows(raw_rows, columns)
235
272
 
273
+ async def query_stream(
274
+ self, sql: str, params: Params = None
275
+ ) -> AsyncGenerator[Row, None]:
276
+ """Yield converted rows on demand; call directly with async for.
277
+
278
+ Submission starts on first iteration. Keep the client open while
279
+ consuming; await aclose() on a retained stream when stopping early.
280
+ """
281
+ _, rows = await self._stream_result(sql, params)
282
+ try:
283
+ async for row in rows:
284
+ yield row
285
+ finally:
286
+ await rows.aclose()
287
+
288
+ async def query_batches(
289
+ self, sql: str, params: Params = None, *, batch_size: int = 1000
290
+ ) -> AsyncGenerator[List[Row], None]:
291
+ """Yield batches of at most batch_size rows across partition boundaries."""
292
+ validate_batch_size(batch_size)
293
+ rows = self.query_stream(sql, params)
294
+ try:
295
+ batch: List[Row] = []
296
+ async for row in rows:
297
+ batch.append(row)
298
+ if len(batch) == batch_size:
299
+ yield batch
300
+ batch = []
301
+ if batch:
302
+ yield batch
303
+ finally:
304
+ await rows.aclose()
305
+
306
+ async def export_csv(self, sql: str, output: TextIO, params: Params = None) -> int:
307
+ """Await query I/O and write CSV to a caller-owned synchronous text sink.
308
+
309
+ Return the data-row count. Failures may leave partial output.
310
+ """
311
+ columns, rows = await self._stream_result(sql, params)
312
+ try:
313
+ writer = CSVWriter(output, columns)
314
+ count = 0
315
+ async for row in rows:
316
+ writer.write(row)
317
+ count += 1
318
+ finally:
319
+ await rows.aclose()
320
+ return count
321
+
322
+ async def export_jsonl(
323
+ self, sql: str, output: TextIO, params: Params = None
324
+ ) -> int:
325
+ """Await query I/O and write JSONL to a caller-owned synchronous text sink.
326
+
327
+ Return the row count. Failures may leave partial output.
328
+ """
329
+ columns, rows = await self._stream_result(sql, params)
330
+ try:
331
+ validate_export_columns(columns)
332
+ count = 0
333
+ async for row in rows:
334
+ write_jsonl(output, row)
335
+ count += 1
336
+ finally:
337
+ await rows.aclose()
338
+ return count
339
+
236
340
  async def query_one(self, sql: str, params: Params = None) -> Optional[Row]:
237
341
  """Return the first row, or ``None`` if the result is empty."""
238
342
  rows = await self.query(sql, params)
@@ -286,6 +390,52 @@ class AsyncSnowflakeClient:
286
390
 
287
391
  # -- async submit -----------------------------------------------------
288
392
 
393
+ async def insert_batches(
394
+ self,
395
+ table: str,
396
+ columns: Sequence[str],
397
+ rows: Iterable[Sequence[Any]],
398
+ *,
399
+ max_rows: int = 1000,
400
+ max_bytes: int = 1_000_000,
401
+ ) -> BatchWriteResult:
402
+ """Insert a single-pass source sequentially within row and wire-byte budgets.
403
+
404
+ Independent statements can commit before a later failure. BatchWriteError
405
+ reports confirmed progress and uncertain outcomes; see docs/batch-writes.md.
406
+ Input is an ordinary iterable, consumed on the calling thread.
407
+ """
408
+ planner = _BatchPlanner(
409
+ table, columns, rows, max_rows, max_bytes, self._build_payload
410
+ )
411
+ while (batch := planner.next_batch()) is not None:
412
+ try:
413
+ self._notify(batch.sql, batch.params)
414
+ except Exception as exc:
415
+ raise planner.error(exc, "validation") from exc
416
+ execution = _BatchExecution(planner, batch)
417
+ try:
418
+ response = await self._transport.submit(
419
+ {}, request_id=execution.failed.request_id, content=batch.content
420
+ )
421
+ execution.received(
422
+ response.statement_handle,
423
+ response.submission_attempts,
424
+ response.is_running,
425
+ )
426
+ if response.is_running:
427
+ handle = response.statement_handle
428
+ if handle is None:
429
+ raise SnowflakeRequestError("running statement has no handle")
430
+ response = await self._poll(handle, self._statement_timeout)
431
+ execution.succeeded = True
432
+ affected = _rows_affected(response.body, strict=True)
433
+ except Exception as exc:
434
+ raise execution.error(exc) from exc
435
+ planner.confirm(batch, affected)
436
+ del batch
437
+ return planner.result()
438
+
289
439
  async def submit(self, sql: str, params: Params = None) -> AsyncQueryHandle:
290
440
  """Submit a long-running statement and return an :class:`AsyncQueryHandle`."""
291
441
  response = await self._run(sql, params, async_exec=True)