Skip to content

Commit e96e5b8

Browse files
large query results
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 716304b commit e96e5b8

File tree

2 files changed

+23
-18
lines changed

2 files changed

+23
-18
lines changed

examples/experimental/tests/test_sea_async_query.py

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -51,12 +51,12 @@ def test_sea_async_query_with_cloud_fetch():
5151
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
5252
)
5353

54-
# Execute a simple query asynchronously
54+
# Execute a query that returns 100 rows asynchronously
5555
cursor = connection.cursor()
56-
logger.info(
57-
"Executing asynchronous query with cloud fetch: SELECT 1 as test_value"
56+
logger.info("Executing asynchronous query with cloud fetch: SELECT 100 rows")
57+
cursor.execute_async(
58+
"SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
5859
)
59-
cursor.execute_async("SELECT 1 as test_value")
6060
logger.info(
6161
"Asynchronous query submitted successfully with cloud fetch enabled"
6262
)
@@ -69,6 +69,8 @@ def test_sea_async_query_with_cloud_fetch():
6969

7070
logger.info("Query is no longer pending, getting results...")
7171
cursor.get_async_execution_result()
72+
rows = cursor.fetchall()
73+
logger.info(f"Retrieved rows: {rows}")
7274
logger.info(
7375
"Successfully retrieved asynchronous query results with cloud fetch enabled"
7476
)
@@ -130,12 +132,12 @@ def test_sea_async_query_without_cloud_fetch():
130132
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
131133
)
132134

133-
# Execute a simple query asynchronously
135+
# Execute a query that returns 100 rows asynchronously
134136
cursor = connection.cursor()
135-
logger.info(
136-
"Executing asynchronous query without cloud fetch: SELECT 1 as test_value"
137+
logger.info("Executing asynchronous query without cloud fetch: SELECT 100 rows")
138+
cursor.execute_async(
139+
"SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
137140
)
138-
cursor.execute_async("SELECT 1 as test_value")
139141
logger.info(
140142
"Asynchronous query submitted successfully with cloud fetch disabled"
141143
)
@@ -148,6 +150,8 @@ def test_sea_async_query_without_cloud_fetch():
148150

149151
logger.info("Query is no longer pending, getting results...")
150152
cursor.get_async_execution_result()
153+
rows = cursor.fetchall()
154+
logger.info(f"Retrieved rows: {rows}")
151155
logger.info(
152156
"Successfully retrieved asynchronous query results with cloud fetch disabled"
153157
)

examples/experimental/tests/test_sea_sync_query.py

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,13 +49,14 @@ def test_sea_sync_query_with_cloud_fetch():
4949
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
5050
)
5151

52-
# Execute a simple query
52+
# Execute a query that returns 100 rows
5353
cursor = connection.cursor()
54-
logger.info(
55-
"Executing synchronous query with cloud fetch: SELECT 1 as test_value"
54+
logger.info("Executing synchronous query with cloud fetch: SELECT 100 rows")
55+
cursor.execute(
56+
"SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
5657
)
57-
cursor.execute("SELECT 1 as test_value")
58-
logger.info("Query executed successfully with cloud fetch enabled")
58+
rows = cursor.fetchall()
59+
logger.info(f"Retrieved rows: {rows}")
5960

6061
# Close resources
6162
cursor.close()
@@ -114,16 +115,16 @@ def test_sea_sync_query_without_cloud_fetch():
114115
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
115116
)
116117

117-
# Execute a simple query
118+
# Execute a query that returns 100 rows
118119
cursor = connection.cursor()
119-
logger.info(
120-
"Executing synchronous query without cloud fetch: SELECT 1 as test_value"
120+
logger.info("Executing synchronous query without cloud fetch: SELECT 100 rows")
121+
cursor.execute(
122+
"SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
121123
)
122-
cursor.execute("SELECT 1 as test_value")
123124
logger.info("Query executed successfully with cloud fetch disabled")
124125

125126
rows = cursor.fetchall()
126-
logger.info(f"Rows: {rows}")
127+
logger.info(f"Retrieved rows: {rows}")
127128

128129
# Close resources
129130
cursor.close()

0 commit comments

Comments
 (0)