Skip to content

Commit 28b4d7b

Browse files
mandatory Manifest, clean up statement_id typing
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 0bdf8f9 commit 28b4d7b

File tree

3 files changed

+15
-5
lines changed

3 files changed

+15
-5
lines changed

src/databricks/sql/backend/sea/utils/filters.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,15 @@ def _filter_sea_result_set(
7373
from databricks.sql.result_set import SeaResultSet
7474

7575
# Create a new SeaResultSet with the filtered data
76+
manifest = result_set.manifest
77+
manifest.total_row_count = len(filtered_rows)
78+
7679
filtered_result_set = SeaResultSet(
7780
connection=result_set.connection,
7881
execute_response=execute_response,
7982
sea_client=cast(SeaDatabricksClient, result_set.backend),
8083
result_data=result_data,
84+
manifest=manifest,
8185
buffer_size_bytes=result_set.buffer_size_bytes,
8286
arraysize=result_set.arraysize,
8387
)

src/databricks/sql/backend/types.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -369,7 +369,7 @@ def from_sea_statement_id(cls, statement_id: str):
369369

370370
return cls(BackendType.SEA, statement_id)
371371

372-
def to_thrift_handle(self):
372+
def to_thrift_handle(self) -> Optional[ttypes.TOperationHandle]:
373373
"""
374374
Convert this CommandId to a Thrift TOperationHandle.
375375
@@ -390,7 +390,7 @@ def to_thrift_handle(self):
390390
modifiedRowCount=self.modified_row_count,
391391
)
392392

393-
def to_sea_statement_id(self):
393+
def to_sea_statement_id(self) -> Optional[str]:
394394
"""
395395
Get the SEA statement ID string.
396396
@@ -401,7 +401,7 @@ def to_sea_statement_id(self):
401401
if self.backend_type != BackendType.SEA:
402402
return None
403403

404-
return self.guid
404+
return str(self.guid)
405405

406406
def to_hex_guid(self) -> str:
407407
"""

src/databricks/sql/result_set.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -473,10 +473,16 @@ def __init__(
473473
manifest: Manifest from SEA response
474474
"""
475475

476+
self.manifest = manifest
477+
478+
statement_id = execute_response.command_id.to_sea_statement_id()
479+
if statement_id is None:
480+
raise ValueError("Command ID is not a SEA statement ID")
481+
476482
results_queue = SeaResultSetQueueFactory.build_queue(
477483
result_data,
478-
manifest,
479-
str(execute_response.command_id.to_sea_statement_id()),
484+
self.manifest,
485+
statement_id,
480486
description=execute_response.description,
481487
max_download_threads=sea_client.max_download_threads,
482488
sea_client=sea_client,

0 commit comments

Comments
 (0)