Skip to content

Commit 5380c7a

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
use more fetch methods
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 015fb76 commit 5380c7a

File tree

2 files changed

+119
-31
lines changed

2 files changed

+119
-31
lines changed

‎examples/experimental/tests/test_sea_async_query.py

Lines changed: 59 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -78,22 +78,44 @@ def test_sea_async_query_with_cloud_fetch():
7878
logger.info("Query is no longer pending, getting results...")
7979
cursor.get_async_execution_result()
8080

81-
# Fetch all rows
82-
rows = cursor.fetchall()
83-
actual_row_count = len(rows)
84-
81+
# Use a mix of fetch methods to retrieve all rows
82+
logger.info("Retrieving data using a mix of fetch methods")
83+
84+
# First, get one row with fetchone
85+
first_row = cursor.fetchone()
86+
if not first_row:
87+
logger.error("FAIL: fetchone returned None, expected a row")
88+
return False
89+
90+
logger.info(f"Successfully retrieved first row with ID: {first_row[0]}")
91+
retrieved_rows = [first_row]
92+
93+
# Then, get a batch of rows with fetchmany
94+
batch_size = 100
95+
batch_rows = cursor.fetchmany(batch_size)
96+
logger.info(f"Successfully retrieved {len(batch_rows)} rows with fetchmany")
97+
retrieved_rows.extend(batch_rows)
98+
99+
# Finally, get all remaining rows with fetchall
100+
remaining_rows = cursor.fetchall()
101+
logger.info(f"Successfully retrieved {len(remaining_rows)} rows with fetchall")
102+
retrieved_rows.extend(remaining_rows)
103+
104+
# Calculate total row count
105+
actual_row_count = len(retrieved_rows)
106+
85107
logger.info(
86108
f"Requested {requested_row_count} rows, received {actual_row_count} rows"
87109
)
88-
89-
# Verify row count
110+
111+
# Verify total row count
90112
if actual_row_count != requested_row_count:
91113
logger.error(
92114
f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
93115
)
94116
return False
95-
96-
logger.info("PASS: Received correct number of rows with cloud fetch")
117+
118+
logger.info("PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly")
97119

98120
# Close resources
99121
cursor.close()
@@ -179,22 +201,44 @@ def test_sea_async_query_without_cloud_fetch():
179201
logger.info("Query is no longer pending, getting results...")
180202
cursor.get_async_execution_result()
181203

182-
# Fetch all rows
183-
rows = cursor.fetchall()
184-
actual_row_count = len(rows)
185-
204+
# Use a mix of fetch methods to retrieve all rows
205+
logger.info("Retrieving data using a mix of fetch methods")
206+
207+
# First, get one row with fetchone
208+
first_row = cursor.fetchone()
209+
if not first_row:
210+
logger.error("FAIL: fetchone returned None, expected a row")
211+
return False
212+
213+
logger.info(f"Successfully retrieved first row with ID: {first_row[0]}")
214+
retrieved_rows = [first_row]
215+
216+
# Then, get a batch of rows with fetchmany
217+
batch_size = 10 # Smaller batch size for non-cloud fetch
218+
batch_rows = cursor.fetchmany(batch_size)
219+
logger.info(f"Successfully retrieved {len(batch_rows)} rows with fetchmany")
220+
retrieved_rows.extend(batch_rows)
221+
222+
# Finally, get all remaining rows with fetchall
223+
remaining_rows = cursor.fetchall()
224+
logger.info(f"Successfully retrieved {len(remaining_rows)} rows with fetchall")
225+
retrieved_rows.extend(remaining_rows)
226+
227+
# Calculate total row count
228+
actual_row_count = len(retrieved_rows)
229+
186230
logger.info(
187231
f"Requested {requested_row_count} rows, received {actual_row_count} rows"
188232
)
189-
190-
# Verify row count
233+
234+
# Verify total row count
191235
if actual_row_count != requested_row_count:
192236
logger.error(
193237
f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
194238
)
195239
return False
196240

197-
logger.info("PASS: Received correct number of rows without cloud fetch")
241+
logger.info("PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly")
198242

199243
# Close resources
200244
cursor.close()

‎examples/experimental/tests/test_sea_sync_query.py

Lines changed: 60 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -64,22 +64,44 @@ def test_sea_sync_query_with_cloud_fetch():
6464
)
6565
cursor.execute(query)
6666

67-
# Fetch all rows
68-
rows = cursor.fetchall()
69-
actual_row_count = len(rows)
70-
67+
# Use a mix of fetch methods to retrieve all rows
68+
logger.info("Retrieving data using a mix of fetch methods")
69+
70+
# First, get one row with fetchone
71+
first_row = cursor.fetchone()
72+
if not first_row:
73+
logger.error("FAIL: fetchone returned None, expected a row")
74+
return False
75+
76+
logger.info(f"Successfully retrieved first row with ID: {first_row[0]}")
77+
retrieved_rows = [first_row]
78+
79+
# Then, get a batch of rows with fetchmany
80+
batch_size = 100
81+
batch_rows = cursor.fetchmany(batch_size)
82+
logger.info(f"Successfully retrieved {len(batch_rows)} rows with fetchmany")
83+
retrieved_rows.extend(batch_rows)
84+
85+
# Finally, get all remaining rows with fetchall
86+
remaining_rows = cursor.fetchall()
87+
logger.info(f"Successfully retrieved {len(remaining_rows)} rows with fetchall")
88+
retrieved_rows.extend(remaining_rows)
89+
90+
# Calculate total row count
91+
actual_row_count = len(retrieved_rows)
92+
7193
logger.info(
7294
f"Requested {requested_row_count} rows, received {actual_row_count} rows"
7395
)
74-
75-
# Verify row count
96+
97+
# Verify total row count
7698
if actual_row_count != requested_row_count:
7799
logger.error(
78100
f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
79101
)
80102
return False
81-
82-
logger.info("PASS: Received correct number of rows with cloud fetch")
103+
104+
logger.info("PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly")
83105

84106
# Close resources
85107
cursor.close()
@@ -153,22 +175,44 @@ def test_sea_sync_query_without_cloud_fetch():
153175
)
154176
cursor.execute(query)
155177

156-
# Fetch all rows
157-
rows = cursor.fetchall()
158-
actual_row_count = len(rows)
159-
178+
# Use a mix of fetch methods to retrieve all rows
179+
logger.info("Retrieving data using a mix of fetch methods")
180+
181+
# First, get one row with fetchone
182+
first_row = cursor.fetchone()
183+
if not first_row:
184+
logger.error("FAIL: fetchone returned None, expected a row")
185+
return False
186+
187+
logger.info(f"Successfully retrieved first row with ID: {first_row[0]}")
188+
retrieved_rows = [first_row]
189+
190+
# Then, get a batch of rows with fetchmany
191+
batch_size = 10 # Smaller batch size for non-cloud fetch
192+
batch_rows = cursor.fetchmany(batch_size)
193+
logger.info(f"Successfully retrieved {len(batch_rows)} rows with fetchmany")
194+
retrieved_rows.extend(batch_rows)
195+
196+
# Finally, get all remaining rows with fetchall
197+
remaining_rows = cursor.fetchall()
198+
logger.info(f"Successfully retrieved {len(remaining_rows)} rows with fetchall")
199+
retrieved_rows.extend(remaining_rows)
200+
201+
# Calculate total row count
202+
actual_row_count = len(retrieved_rows)
203+
160204
logger.info(
161205
f"Requested {requested_row_count} rows, received {actual_row_count} rows"
162206
)
163-
164-
# Verify row count
207+
208+
# Verify total row count
165209
if actual_row_count != requested_row_count:
166210
logger.error(
167211
f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
168212
)
169213
return False
170-
171-
logger.info("PASS: Received correct number of rows without cloud fetch")
214+
215+
logger.info("PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly")
172216

173217
# Close resources
174218
cursor.close()

0 commit comments

Comments
 (0)