Skip to content

Commit 91d28b2

Browse files
remove parsing in backend
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 68d6276 commit 91d28b2

File tree

1 file changed

+14
-22
lines changed

1 file changed

+14
-22
lines changed

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

Lines changed: 14 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,9 @@ def _extract_description_from_manifest(
323323

324324
return columns if columns else None
325325

326-
def _results_message_to_execute_response(self, sea_response, command_id):
326+
def _results_message_to_execute_response(
327+
self, response: GetStatementResponse
328+
) -> ExecuteResponse:
327329
"""
328330
Convert a SEA response to an ExecuteResponse and extract result data.
329331
@@ -332,33 +334,27 @@ def _results_message_to_execute_response(self, sea_response, command_id):
332334
command_id: The command ID
333335
334336
Returns:
335-
tuple: (ExecuteResponse, ResultData, ResultManifest) - The normalized execute response,
336-
result data object, and manifest object
337+
ExecuteResponse: The normalized execute response
337338
"""
338339

339-
# Parse the response
340-
status = parse_status(sea_response)
341-
manifest_obj = parse_manifest(sea_response)
342-
result_data_obj = parse_result(sea_response)
343-
344340
# Extract description from manifest schema
345-
description = self._extract_description_from_manifest(manifest_obj)
341+
description = self._extract_description_from_manifest(response.manifest)
346342

347343
# Check for compression
348-
lz4_compressed = manifest_obj.result_compression == "LZ4_FRAME"
344+
lz4_compressed = response.manifest.result_compression == "LZ4_FRAME"
349345

350346
execute_response = ExecuteResponse(
351-
command_id=command_id,
352-
status=status.state,
347+
command_id=response.statement_id,
348+
status=response.status.state,
353349
description=description,
354350
has_been_closed_server_side=False,
355351
lz4_compressed=lz4_compressed,
356352
is_staging_operation=False,
357353
arrow_schema_bytes=None,
358-
result_format=manifest_obj.format,
354+
result_format=response.manifest.format,
359355
)
360356

361-
return execute_response, result_data_obj, manifest_obj
357+
return execute_response
362358

363359
def execute_command(
364360
self,
@@ -586,25 +582,21 @@ def get_execution_result(
586582
path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
587583
data=request.to_dict(),
588584
)
585+
response = GetStatementResponse.from_dict(response_data)
589586

590587
# Create and return a SeaResultSet
591588
from databricks.sql.result_set import SeaResultSet
592589

593-
# Convert the response to an ExecuteResponse and extract result data
594-
(
595-
execute_response,
596-
result_data,
597-
manifest,
598-
) = self._results_message_to_execute_response(response_data, command_id)
590+
execute_response = self._results_message_to_execute_response(response)
599591

600592
return SeaResultSet(
601593
connection=cursor.connection,
602594
execute_response=execute_response,
603595
sea_client=self,
604596
buffer_size_bytes=cursor.buffer_size_bytes,
605597
arraysize=cursor.arraysize,
606-
result_data=result_data,
607-
manifest=manifest,
598+
result_data=response.result,
599+
manifest=response.manifest,
608600
)
609601

610602
# == Metadata Operations ==

0 commit comments

Comments
 (0)