|
21 | 21 |
|
22 | 22 | from __future__ import annotations |
23 | 23 |
|
| 24 | +import hashlib |
24 | 25 | import sys |
25 | 26 | from uuid import uuid4 |
26 | 27 |
|
@@ -136,6 +137,33 @@ def conn(kernel_conn_params): |
136 | 137 | c.close() |
137 | 138 |
|
138 | 139 |
|
| 140 | +def test_cloudfetch_disabled_returns_complete_large_result(kernel_conn_params): |
| 141 | + from databricks.sql.backend.kernel.client import KernelDatabricksClient |
| 142 | + |
| 143 | + with sql.connect(**kernel_conn_params, use_cloud_fetch=False) as connection: |
| 144 | + assert isinstance(connection.session.backend, KernelDatabricksClient) |
| 145 | + with connection.cursor() as cursor: |
| 146 | + cursor.execute( |
| 147 | + "SELECT id, sha2(CAST(id AS STRING), 512) AS payload " |
| 148 | + "FROM range(250000) ORDER BY id" |
| 149 | + ) |
| 150 | + count = 0 |
| 151 | + payload_bytes = 0 |
| 152 | + digest = hashlib.sha256() |
| 153 | + while rows := cursor.fetchmany(8192): |
| 154 | + for identity, payload in rows: |
| 155 | + assert identity == count |
| 156 | + assert payload == hashlib.sha512(str(identity).encode()).hexdigest() |
| 157 | + payload_bytes += len(payload) |
| 158 | + digest.update(f"{identity}\t{payload}\n".encode()) |
| 159 | + count += 1 |
| 160 | + assert count == 250000 |
| 161 | + assert payload_bytes == 32000000 |
| 162 | + assert digest.hexdigest() == ( |
| 163 | + "a25ada4c6c476b53e97eb6e8a945e9ce4c14cc432c7ad672dc6fb607d660c86a" |
| 164 | + ) |
| 165 | + |
| 166 | + |
139 | 167 | def test_connect_with_use_kernel_opens_a_session(conn): |
140 | 168 | assert conn.open, "connection should report open after connect()" |
141 | 169 |
|
@@ -366,9 +394,7 @@ def test_schemas_with_empty_string_filter_matches_nothing(conn): |
366 | 394 | assert cur.fetchall() == [] |
367 | 395 |
|
368 | 396 |
|
369 | | -@pytest.mark.parametrize( |
370 | | - "empty_filter", ["schema_name", "table_name", "column_name"] |
371 | | -) |
| 397 | +@pytest.mark.parametrize("empty_filter", ["schema_name", "table_name", "column_name"]) |
372 | 398 | def test_columns_with_empty_string_filter_matches_nothing(conn, empty_filter): |
373 | 399 | filters = { |
374 | 400 | "catalog_name": "system", |
|
0 commit comments