Skip to content

Commit c038d5a

Browse files
working JSON + INLINE
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 40f79b5 commit c038d5a

File tree

3 files changed

+45
-19
lines changed

3 files changed

+45
-19
lines changed

examples/experimental/tests/test_sea_metadata.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,26 +56,34 @@ def test_sea_metadata():
5656
cursor = connection.cursor()
5757
logger.info("Fetching catalogs...")
5858
cursor.catalogs()
59+
rows = cursor.fetchall()
60+
logger.info(f"Rows: {rows}")
5961
logger.info("Successfully fetched catalogs")
6062

6163
# Test schemas
6264
logger.info(f"Fetching schemas for catalog '{catalog}'...")
6365
cursor.schemas(catalog_name=catalog)
66+
rows = cursor.fetchall()
67+
logger.info(f"Rows: {rows}")
6468
logger.info("Successfully fetched schemas")
6569

6670
# Test tables
6771
logger.info(f"Fetching tables for catalog '{catalog}', schema 'default'...")
6872
cursor.tables(catalog_name=catalog, schema_name="default")
73+
rows = cursor.fetchall()
74+
logger.info(f"Rows: {rows}")
6975
logger.info("Successfully fetched tables")
7076

7177
# Test columns for a specific table
7278
# Using a common table that should exist in most environments
7379
logger.info(
74-
f"Fetching columns for catalog '{catalog}', schema 'default', table 'information_schema'..."
80+
f"Fetching columns for catalog '{catalog}', schema 'default', table 'customer'..."
7581
)
7682
cursor.columns(
77-
catalog_name=catalog, schema_name="default", table_name="information_schema"
83+
catalog_name=catalog, schema_name="default", table_name="customer"
7884
)
85+
rows = cursor.fetchall()
86+
logger.info(f"Rows: {rows}")
7987
logger.info("Successfully fetched columns")
8088

8189
# Close resources

src/databricks/sql/result_set.py

Lines changed: 30 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,12 @@
2020
from databricks.sql.thrift_api.TCLIService import ttypes
2121
from databricks.sql.types import Row
2222
from databricks.sql.exc import Error, RequestError, CursorAlreadyClosedError
23-
from databricks.sql.utils import ColumnTable, ColumnQueue, JsonQueue, SeaResultSetQueueFactory
23+
from databricks.sql.utils import (
24+
ColumnTable,
25+
ColumnQueue,
26+
JsonQueue,
27+
SeaResultSetQueueFactory,
28+
)
2429
from databricks.sql.backend.types import CommandId, CommandState, ExecuteResponse
2530

2631
logger = logging.getLogger(__name__)
@@ -469,8 +474,8 @@ def __init__(
469474
sea_client: "SeaDatabricksClient",
470475
buffer_size_bytes: int = 104857600,
471476
arraysize: int = 10000,
472-
result_data=None,
473-
manifest=None,
477+
result_data: Optional[ResultData] = None,
478+
manifest: Optional[ResultManifest] = None,
474479
):
475480
"""
476481
Initialize a SeaResultSet with the response from a SEA query execution.
@@ -485,13 +490,17 @@ def __init__(
485490
manifest: Manifest from SEA response (optional)
486491
"""
487492

488-
queue = SeaResultSetQueueFactory.build_queue(
489-
sea_result_data=execute_response.results_data,
490-
manifest=execute_response.results_manifest,
491-
statement_id=execute_response.command_id.to_sea_statement_id(),
492-
description=execute_response.description,
493-
schema_bytes=execute_response.arrow_schema_bytes,
494-
)
493+
if result_data:
494+
queue = SeaResultSetQueueFactory.build_queue(
495+
sea_result_data=result_data,
496+
manifest=manifest,
497+
statement_id=execute_response.command_id.to_sea_statement_id(),
498+
description=execute_response.description,
499+
schema_bytes=execute_response.arrow_schema_bytes,
500+
)
501+
else:
502+
logger.warning("No result data provided for SEA result set")
503+
queue = JsonQueue([])
495504

496505
super().__init__(
497506
connection=connection,
@@ -501,12 +510,13 @@ def __init__(
501510
command_id=execute_response.command_id,
502511
status=execute_response.status,
503512
has_been_closed_server_side=execute_response.has_been_closed_server_side,
513+
results_queue=queue,
504514
description=execute_response.description,
505515
is_staging_operation=execute_response.is_staging_operation,
506516
lz4_compressed=execute_response.lz4_compressed,
507517
arrow_schema_bytes=execute_response.arrow_schema_bytes,
508518
)
509-
519+
510520
def _convert_to_row_objects(self, rows):
511521
"""
512522
Convert raw data rows to Row objects with named columns based on description.
@@ -526,9 +536,7 @@ def _convert_to_row_objects(self, rows):
526536

527537
def _fill_results_buffer(self):
528538
"""Fill the results buffer from the backend."""
529-
raise NotImplementedError(
530-
"_fill_results_buffer is not implemented for SEA backend"
531-
)
539+
return None
532540

533541
def fetchone(self) -> Optional[Row]:
534542
"""
@@ -572,8 +580,15 @@ def fetchall(self) -> List[Row]:
572580
"""
573581
Fetch all (remaining) rows of a query result, returning them as a list of rows.
574582
"""
583+
# Note: We check for the specific queue type to maintain consistency with ThriftResultSet
584+
if isinstance(self.results, JsonQueue):
585+
rows = self.results.remaining_rows()
586+
self._next_row_index += len(rows)
575587

576-
raise NotImplementedError("fetchall is not implemented for SEA backend")
588+
# Convert to Row objects
589+
return self._convert_to_row_objects(rows)
590+
else:
591+
raise NotImplementedError("Unsupported queue type")
577592

578593
def fetchmany_arrow(self, size: int) -> Any:
579594
"""Fetch the next set of rows as an Arrow table."""
@@ -606,4 +621,3 @@ def fetchall_arrow(self) -> Any:
606621
return self._convert_rows_to_arrow_table(rows)
607622
else:
608623
raise NotImplementedError("Unsupported queue type")
609-

src/databricks/sql/utils.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ def build_queue(
109109
else:
110110
raise AssertionError("Row set type is not valid")
111111

112+
112113
class SeaResultSetQueueFactory(ABC):
113114
@staticmethod
114115
def build_queue(
@@ -145,7 +146,9 @@ def build_queue(
145146
return JsonQueue(sea_result_data.data)
146147
elif sea_result_data.external_links is not None:
147148
# EXTERNAL_LINKS disposition
148-
raise NotImplementedError("EXTERNAL_LINKS disposition is not implemented for SEA backend")
149+
raise NotImplementedError(
150+
"EXTERNAL_LINKS disposition is not implemented for SEA backend"
151+
)
149152
else:
150153
# Empty result set
151154
return JsonQueue([])
@@ -173,6 +176,7 @@ def remaining_rows(self):
173176
self.cur_row_index += len(slice)
174177
return slice
175178

179+
176180
class ColumnTable:
177181
def __init__(self, column_table, column_names):
178182
self.column_table = column_table

0 commit comments

Comments
 (0)