Skip to content

Commit e43c07b

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
Revert "Complete Fetch Phase (for INLINE disposition and JSON_ARRAY format) (#594)"
This reverts commit 70c7dc8.
1 parent ab2e43d commit e43c07b

19 files changed

Lines changed: 261 additions & 1394 deletions

‎examples/experimental/tests/test_sea_async_query.py‎

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

55-
# Execute a query that generates large rows to force multiple chunks
56-
requested_row_count = 5000
55+
# Execute a simple query asynchronously
5756
cursor = connection.cursor()
58-
query = f"""
59-
SELECT
60-
id,
61-
concat('value_', repeat('a', 10000)) as test_value
62-
FROM range(1, {requested_row_count} + 1) AS t(id)
63-
"""
64-
6557
logger.info(
66-
f"Executing asynchronous query with cloud fetch to generate {requested_row_count} rows"
58+
"Executing asynchronous query with cloud fetch: SELECT 1 as test_value"
6759
)
68-
cursor.execute_async(query)
60+
cursor.execute_async("SELECT 1 as test_value")
6961
logger.info(
7062
"Asynchronous query submitted successfully with cloud fetch enabled"
7163
)
@@ -78,25 +70,8 @@ def test_sea_async_query_with_cloud_fetch():
7870

7971
logger.info("Query is no longer pending, getting results...")
8072
cursor.get_async_execution_result()
81-
82-
results = [cursor.fetchone()]
83-
results.extend(cursor.fetchmany(10))
84-
results.extend(cursor.fetchall())
85-
actual_row_count = len(results)
86-
87-
logger.info(
88-
f"Requested {requested_row_count} rows, received {actual_row_count} rows"
89-
)
90-
91-
# Verify total row count
92-
if actual_row_count != requested_row_count:
93-
logger.error(
94-
f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
95-
)
96-
return False
97-
9873
logger.info(
99-
"PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly"
74+
"Successfully retrieved asynchronous query results with cloud fetch enabled"
10075
)
10176

10277
# Close resources
@@ -156,20 +131,12 @@ def test_sea_async_query_without_cloud_fetch():
156131
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
157132
)
158133

159-
# For non-cloud fetch, use a smaller row count to avoid exceeding inline limits
160-
requested_row_count = 100
134+
# Execute a simple query asynchronously
161135
cursor = connection.cursor()
162-
query = f"""
163-
SELECT
164-
id,
165-
concat('value_', repeat('a', 100)) as test_value
166-
FROM range(1, {requested_row_count} + 1) AS t(id)
167-
"""
168-
169136
logger.info(
170-
f"Executing asynchronous query without cloud fetch to generate {requested_row_count} rows"
137+
"Executing asynchronous query without cloud fetch: SELECT 1 as test_value"
171138
)
172-
cursor.execute_async(query)
139+
cursor.execute_async("SELECT 1 as test_value")
173140
logger.info(
174141
"Asynchronous query submitted successfully with cloud fetch disabled"
175142
)
@@ -182,24 +149,8 @@ def test_sea_async_query_without_cloud_fetch():
182149

183150
logger.info("Query is no longer pending, getting results...")
184151
cursor.get_async_execution_result()
185-
results = [cursor.fetchone()]
186-
results.extend(cursor.fetchmany(10))
187-
results.extend(cursor.fetchall())
188-
actual_row_count = len(results)
189-
190-
logger.info(
191-
f"Requested {requested_row_count} rows, received {actual_row_count} rows"
192-
)
193-
194-
# Verify total row count
195-
if actual_row_count != requested_row_count:
196-
logger.error(
197-
f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
198-
)
199-
return False
200-
201152
logger.info(
202-
"PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly"
153+
"Successfully retrieved asynchronous query results with cloud fetch disabled"
203154
)
204155

205156
# Close resources

‎examples/experimental/tests/test_sea_sync_query.py‎

Lines changed: 9 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -50,27 +50,13 @@ def test_sea_sync_query_with_cloud_fetch():
5050
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
5151
)
5252

53-
# Execute a query that generates large rows to force multiple chunks
54-
requested_row_count = 10000
53+
# Execute a simple query
5554
cursor = connection.cursor()
56-
query = f"""
57-
SELECT
58-
id,
59-
concat('value_', repeat('a', 10000)) as test_value
60-
FROM range(1, {requested_row_count} + 1) AS t(id)
61-
"""
62-
63-
logger.info(
64-
f"Executing synchronous query with cloud fetch to generate {requested_row_count} rows"
65-
)
66-
cursor.execute(query)
67-
results = [cursor.fetchone()]
68-
results.extend(cursor.fetchmany(10))
69-
results.extend(cursor.fetchall())
70-
actual_row_count = len(results)
7155
logger.info(
72-
f"{actual_row_count} rows retrieved against {requested_row_count} requested"
56+
"Executing synchronous query with cloud fetch: SELECT 1 as test_value"
7357
)
58+
cursor.execute("SELECT 1 as test_value")
59+
logger.info("Query executed successfully with cloud fetch enabled")
7460

7561
# Close resources
7662
cursor.close()
@@ -129,18 +115,13 @@ def test_sea_sync_query_without_cloud_fetch():
129115
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
130116
)
131117

132-
# For non-cloud fetch, use a smaller row count to avoid exceeding inline limits
133-
requested_row_count = 100
118+
# Execute a simple query
134119
cursor = connection.cursor()
135-
logger.info("Executing synchronous query without cloud fetch: SELECT 100 rows")
136-
cursor.execute(
137-
"SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
120+
logger.info(
121+
"Executing synchronous query without cloud fetch: SELECT 1 as test_value"
138122
)
139-
140-
results = [cursor.fetchone()]
141-
results.extend(cursor.fetchmany(10))
142-
results.extend(cursor.fetchall())
143-
logger.info(f"{len(results)} rows retrieved against 100 requested")
123+
cursor.execute("SELECT 1 as test_value")
124+
logger.info("Query executed successfully with cloud fetch disabled")
144125

145126
# Close resources
146127
cursor.close()

‎src/databricks/sql/backend/sea/backend.py‎

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,7 @@ def close_session(self, session_id: SessionId) -> None:
259259
logger.debug("SeaDatabricksClient.close_session(session_id=%s)", session_id)
260260

261261
if session_id.backend_type != BackendType.SEA:
262-
raise ValueError("Not a valid SEA session ID")
262+
raise ProgrammingError("Not a valid SEA session ID")
263263
sea_session_id = session_id.to_sea_session_id()
264264

265265
request_data = DeleteSessionRequest(
@@ -298,7 +298,7 @@ def get_allowed_session_configurations() -> List[str]:
298298

299299
def _extract_description_from_manifest(
300300
self, manifest: ResultManifest
301-
) -> List[Tuple]:
301+
) -> Optional[List]:
302302
"""
303303
Extract column description from a manifest object, in the format defined by
304304
the spec: https://peps.python.org/pep-0249/#description
@@ -307,12 +307,15 @@ def _extract_description_from_manifest(
307307
manifest: The ResultManifest object containing schema information
308308
309309
Returns:
310-
List[Tuple]: A list of column tuples
310+
Optional[List]: A list of column tuples or None if no columns are found
311311
"""
312312

313313
schema_data = manifest.schema
314314
columns_data = schema_data.get("columns", [])
315315

316+
if not columns_data:
317+
return None
318+
316319
columns = []
317320
for col_data in columns_data:
318321
# Format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
@@ -328,7 +331,7 @@ def _extract_description_from_manifest(
328331
)
329332
)
330333

331-
return columns
334+
return columns if columns else None
332335

333336
def _results_message_to_execute_response(
334337
self, response: Union[ExecuteStatementResponse, GetStatementResponse]
@@ -459,7 +462,7 @@ def execute_command(
459462
"""
460463

461464
if session_id.backend_type != BackendType.SEA:
462-
raise ValueError("Not a valid SEA session ID")
465+
raise ProgrammingError("Not a valid SEA session ID")
463466

464467
sea_session_id = session_id.to_sea_session_id()
465468

@@ -547,11 +550,9 @@ def cancel_command(self, command_id: CommandId) -> None:
547550
"""
548551

549552
if command_id.backend_type != BackendType.SEA:
550-
raise ValueError("Not a valid SEA command ID")
553+
raise ProgrammingError("Not a valid SEA command ID")
551554

552555
sea_statement_id = command_id.to_sea_statement_id()
553-
if sea_statement_id is None:
554-
raise ValueError("Not a valid SEA command ID")
555556

556557
request = CancelStatementRequest(statement_id=sea_statement_id)
557558
self.http_client._make_request(
@@ -572,11 +573,9 @@ def close_command(self, command_id: CommandId) -> None:
572573
"""
573574

574575
if command_id.backend_type != BackendType.SEA:
575-
raise ValueError("Not a valid SEA command ID")
576+
raise ProgrammingError("Not a valid SEA command ID")
576577

577578
sea_statement_id = command_id.to_sea_statement_id()
578-
if sea_statement_id is None:
579-
raise ValueError("Not a valid SEA command ID")
580579

581580
request = CloseStatementRequest(statement_id=sea_statement_id)
582581
self.http_client._make_request(
@@ -594,8 +593,6 @@ def _poll_query(self, command_id: CommandId) -> GetStatementResponse:
594593
raise ValueError("Not a valid SEA command ID")
595594

596595
sea_statement_id = command_id.to_sea_statement_id()
597-
if sea_statement_id is None:
598-
raise ValueError("Not a valid SEA command ID")
599596

600597
request = GetStatementRequest(statement_id=sea_statement_id)
601598
response_data = self.http_client._make_request(

‎src/databricks/sql/backend/sea/queue.py‎

Lines changed: 0 additions & 86 deletions
This file was deleted.

0 commit comments

Comments
 (0)