Skip to content

Commit eece87e

Browse files
test(kernel): fetch complete large results with CloudFetch disabled
Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
1 parent 3b57c66 commit eece87e

1 file changed

Lines changed: 29 additions & 3 deletions

File tree

‎tests/e2e/test_kernel_backend.py‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
from __future__ import annotations
2323

24+
import hashlib
2425
import sys
2526
from uuid import uuid4
2627

@@ -136,6 +137,33 @@ def conn(kernel_conn_params):
136137
c.close()
137138

138139

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+
139167
def test_connect_with_use_kernel_opens_a_session(conn):
140168
assert conn.open, "connection should report open after connect()"
141169

@@ -366,9 +394,7 @@ def test_schemas_with_empty_string_filter_matches_nothing(conn):
366394
assert cur.fetchall() == []
367395

368396

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"])
372398
def test_columns_with_empty_string_filter_matches_nothing(conn, empty_filter):
373399
filters = {
374400
"catalog_name": "system",

0 commit comments

Comments
 (0)