Skip to content

Commit 165c4f3

Browse files
remove un-necessary changes
covered by #588 Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 787f1f7 commit 165c4f3

File tree

4 files changed

+7
-309
lines changed

4 files changed

+7
-309
lines changed

src/databricks/sql/backend/thrift_backend.py

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
CommandId,
1616
ExecuteResponse,
1717
)
18+
from databricks.sql.backend.utils import guid_to_hex_id
1819

1920

2021
try:
@@ -841,8 +842,6 @@ def get_execution_result(
841842

842843
status = self.get_query_state(command_id)
843844

844-
status = self.get_query_state(command_id)
845-
846845
execute_response = ExecuteResponse(
847846
command_id=command_id,
848847
status=status,
@@ -895,7 +894,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
895894
self._check_command_not_in_error_or_closed_state(thrift_handle, poll_resp)
896895
state = CommandState.from_thrift_state(operation_state)
897896
if state is None:
898-
raise ValueError(f"Invalid operation state: {operation_state}")
897+
raise ValueError(f"Unknown command state: {operation_state}")
899898
return state
900899

901900
@staticmethod
@@ -1189,11 +1188,7 @@ def _handle_execute_response(self, resp, cursor):
11891188
resp.directResults and resp.directResults.operationStatus,
11901189
)
11911190

1192-
(
1193-
execute_response,
1194-
arrow_schema_bytes,
1195-
) = self._results_message_to_execute_response(resp, final_operation_state)
1196-
return execute_response, arrow_schema_bytes
1191+
return self._results_message_to_execute_response(resp, final_operation_state)
11971192

11981193
def _handle_execute_response_async(self, resp, cursor):
11991194
command_id = CommandId.from_thrift_handle(resp.operationHandle)

src/databricks/sql/result_set.py

Lines changed: 0 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -438,76 +438,3 @@ def map_col_type(type_):
438438
(column.name, map_col_type(column.datatype), None, None, None, None, None)
439439
for column in table_schema_message.columns
440440
]
441-
442-
443-
class SeaResultSet(ResultSet):
444-
"""ResultSet implementation for the SEA backend."""
445-
446-
def __init__(
447-
self,
448-
connection: "Connection",
449-
execute_response: "ExecuteResponse",
450-
sea_client: "SeaDatabricksClient",
451-
buffer_size_bytes: int = 104857600,
452-
arraysize: int = 10000,
453-
):
454-
"""
455-
Initialize a SeaResultSet with the response from a SEA query execution.
456-
457-
Args:
458-
connection: The parent connection
459-
sea_client: The SeaDatabricksClient instance for direct access
460-
buffer_size_bytes: Buffer size for fetching results
461-
arraysize: Default number of rows to fetch
462-
execute_response: Response from the execute command (new style)
463-
sea_response: Direct SEA response (legacy style)
464-
"""
465-
466-
super().__init__(
467-
connection=connection,
468-
backend=sea_client,
469-
arraysize=arraysize,
470-
buffer_size_bytes=buffer_size_bytes,
471-
command_id=execute_response.command_id,
472-
status=execute_response.status,
473-
has_been_closed_server_side=execute_response.has_been_closed_server_side,
474-
has_more_rows=execute_response.has_more_rows,
475-
results_queue=execute_response.results_queue,
476-
description=execute_response.description,
477-
is_staging_operation=execute_response.is_staging_operation,
478-
)
479-
480-
def _fill_results_buffer(self):
481-
"""Fill the results buffer from the backend."""
482-
raise NotImplementedError("fetchone is not implemented for SEA backend")
483-
484-
def fetchone(self) -> Optional[Row]:
485-
"""
486-
Fetch the next row of a query result set, returning a single sequence,
487-
or None when no more data is available.
488-
"""
489-
490-
raise NotImplementedError("fetchone is not implemented for SEA backend")
491-
492-
def fetchmany(self, size: Optional[int] = None) -> List[Row]:
493-
"""
494-
Fetch the next set of rows of a query result, returning a list of rows.
495-
496-
An empty sequence is returned when no more rows are available.
497-
"""
498-
499-
raise NotImplementedError("fetchmany is not implemented for SEA backend")
500-
501-
def fetchall(self) -> List[Row]:
502-
"""
503-
Fetch all (remaining) rows of a query result, returning them as a list of rows.
504-
"""
505-
raise NotImplementedError("fetchall is not implemented for SEA backend")
506-
507-
def fetchmany_arrow(self, size: int) -> Any:
508-
"""Fetch the next set of rows as an Arrow table."""
509-
raise NotImplementedError("fetchmany_arrow is not implemented for SEA backend")
510-
511-
def fetchall_arrow(self) -> Any:
512-
"""Fetch all remaining rows as an Arrow table."""
513-
raise NotImplementedError("fetchall_arrow is not implemented for SEA backend")

tests/unit/test_sea_result_set.py

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

tests/unit/test_thrift_backend.py

Lines changed: 4 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -619,13 +619,6 @@ def test_handle_execute_response_sets_compression_in_direct_results(
619619
lz4Compressed = Mock()
620620
resultSet = MagicMock()
621621
resultSet.results.startRowOffset = 0
622-
623-
# Create a valid operation status
624-
op_status = ttypes.TGetOperationStatusResp(
625-
status=self.okay_status,
626-
operationState=ttypes.TOperationState.FINISHED_STATE,
627-
)
628-
629622
t_execute_resp = resp_type(
630623
status=Mock(),
631624
operationHandle=Mock(),
@@ -927,9 +920,7 @@ def test_handle_execute_response_can_handle_with_direct_results(self):
927920
auth_provider=AuthProvider(),
928921
ssl_options=SSLOptions(),
929922
)
930-
thrift_backend._results_message_to_execute_response = Mock(
931-
return_value=(Mock(), Mock())
932-
)
923+
thrift_backend._results_message_to_execute_response = Mock()
933924

934925
thrift_backend._handle_execute_response(execute_resp, Mock())
935926

@@ -957,12 +948,6 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
957948
operationHandle=self.operation_handle,
958949
)
959950

960-
# Mock the operation status response
961-
op_state = ttypes.TGetOperationStatusResp(
962-
status=self.okay_status,
963-
operationState=ttypes.TOperationState.FINISHED_STATE,
964-
)
965-
tcli_service_instance.GetOperationStatus.return_value = op_state
966951
tcli_service_instance.GetResultSetMetadata.return_value = (
967952
t_get_result_set_metadata_resp
968953
)
@@ -977,7 +962,7 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
977962
t_execute_resp, Mock()
978963
)
979964

980-
self.assertEqual(arrow_schema_bytes, arrow_schema_mock)
965+
self.assertEqual(execute_response.arrow_schema_bytes, arrow_schema_mock)
981966

982967
@patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
983968
def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
@@ -997,12 +982,6 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
997982
operationHandle=self.operation_handle,
998983
)
999984

1000-
# Mock the operation status response
1001-
op_state = ttypes.TGetOperationStatusResp(
1002-
status=self.okay_status,
1003-
operationState=ttypes.TOperationState.FINISHED_STATE,
1004-
)
1005-
tcli_service_instance.GetOperationStatus.return_value = op_state
1006985
tcli_service_instance.GetResultSetMetadata.return_value = hive_schema_req
1007986
tcli_service_instance.GetOperationStatus.return_value = (
1008987
ttypes.TGetOperationStatusResp(
@@ -1694,9 +1673,7 @@ def test_handle_execute_response_sets_active_op_handle(self):
16941673
thrift_backend = self._make_fake_thrift_backend()
16951674
thrift_backend._check_direct_results_for_error = Mock()
16961675
thrift_backend._wait_until_command_done = Mock()
1697-
thrift_backend._results_message_to_execute_response = Mock(
1698-
return_value=(Mock(), Mock())
1699-
)
1676+
thrift_backend._results_message_to_execute_response = Mock()
17001677

17011678
# Create a mock response with a real operation handle
17021679
mock_resp = Mock()
@@ -2256,8 +2233,7 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
22562233
@patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
22572234
@patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
22582235
@patch(
2259-
"databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response",
2260-
return_value=(Mock(), Mock()),
2236+
"databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response"
22612237
)
22622238
def test_execute_command_sets_complex_type_fields_correctly(
22632239
self, mock_handle_execute_response, tcli_service_class, mock_result_set

0 commit comments

Comments
 (0)