remove un-necessary changes

covered by #588

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -15,6 +15,7 @@
     CommandId,
     ExecuteResponse,
 )
+from databricks.sql.backend.utils import guid_to_hex_id
 
 
 try:
@@ -841,8 +842,6 @@ def get_execution_result(
 
         status = self.get_query_state(command_id)
 
-        status = self.get_query_state(command_id)
-
         execute_response = ExecuteResponse(
             command_id=command_id,
             status=status,
@@ -895,7 +894,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         self._check_command_not_in_error_or_closed_state(thrift_handle, poll_resp)
         state = CommandState.from_thrift_state(operation_state)
         if state is None:
-            raise ValueError(f"Invalid operation state: {operation_state}")
+            raise ValueError(f"Unknown command state: {operation_state}")
         return state
 
     @staticmethod
@@ -1189,11 +1188,7 @@ def _handle_execute_response(self, resp, cursor):
             resp.directResults and resp.directResults.operationStatus,
         )
 
-        (
-            execute_response,
-            arrow_schema_bytes,
-        ) = self._results_message_to_execute_response(resp, final_operation_state)
-        return execute_response, arrow_schema_bytes
+        return self._results_message_to_execute_response(resp, final_operation_state)
 
     def _handle_execute_response_async(self, resp, cursor):
         command_id = CommandId.from_thrift_handle(resp.operationHandle)

src/databricks/sql/result_set.py

@@ -438,76 +438,3 @@ def map_col_type(type_):
             (column.name, map_col_type(column.datatype), None, None, None, None, None)
             for column in table_schema_message.columns
         ]
-
-
-class SeaResultSet(ResultSet):
-    """ResultSet implementation for the SEA backend."""
-
-    def __init__(
-        self,
-        connection: "Connection",
-        execute_response: "ExecuteResponse",
-        sea_client: "SeaDatabricksClient",
-        buffer_size_bytes: int = 104857600,
-        arraysize: int = 10000,
-    ):
-        """
-        Initialize a SeaResultSet with the response from a SEA query execution.
-
-        Args:
-            connection: The parent connection
-            sea_client: The SeaDatabricksClient instance for direct access
-            buffer_size_bytes: Buffer size for fetching results
-            arraysize: Default number of rows to fetch
-            execute_response: Response from the execute command (new style)
-            sea_response: Direct SEA response (legacy style)
-        """
-
-        super().__init__(
-            connection=connection,
-            backend=sea_client,
-            arraysize=arraysize,
-            buffer_size_bytes=buffer_size_bytes,
-            command_id=execute_response.command_id,
-            status=execute_response.status,
-            has_been_closed_server_side=execute_response.has_been_closed_server_side,
-            has_more_rows=execute_response.has_more_rows,
-            results_queue=execute_response.results_queue,
-            description=execute_response.description,
-            is_staging_operation=execute_response.is_staging_operation,
-        )
-
-    def _fill_results_buffer(self):
-        """Fill the results buffer from the backend."""
-        raise NotImplementedError("fetchone is not implemented for SEA backend")
-
-    def fetchone(self) -> Optional[Row]:
-        """
-        Fetch the next row of a query result set, returning a single sequence,
-        or None when no more data is available.
-        """
-
-        raise NotImplementedError("fetchone is not implemented for SEA backend")
-
-    def fetchmany(self, size: Optional[int] = None) -> List[Row]:
-        """
-        Fetch the next set of rows of a query result, returning a list of rows.
-
-        An empty sequence is returned when no more rows are available.
-        """
-
-        raise NotImplementedError("fetchmany is not implemented for SEA backend")
-
-    def fetchall(self) -> List[Row]:
-        """
-        Fetch all (remaining) rows of a query result, returning them as a list of rows.
-        """
-        raise NotImplementedError("fetchall is not implemented for SEA backend")
-
-    def fetchmany_arrow(self, size: int) -> Any:
-        """Fetch the next set of rows as an Arrow table."""
-        raise NotImplementedError("fetchmany_arrow is not implemented for SEA backend")
-
-    def fetchall_arrow(self) -> Any:
-        """Fetch all remaining rows as an Arrow table."""
-        raise NotImplementedError("fetchall_arrow is not implemented for SEA backend")

tests/unit/test_sea_result_set.py

@@ -619,13 +619,6 @@ def test_handle_execute_response_sets_compression_in_direct_results(
             lz4Compressed = Mock()
             resultSet = MagicMock()
             resultSet.results.startRowOffset = 0
-
-            # Create a valid operation status
-            op_status = ttypes.TGetOperationStatusResp(
-                status=self.okay_status,
-                operationState=ttypes.TOperationState.FINISHED_STATE,
-            )
-
             t_execute_resp = resp_type(
                 status=Mock(),
                 operationHandle=Mock(),
@@ -927,9 +920,7 @@ def test_handle_execute_response_can_handle_with_direct_results(self):
                     auth_provider=AuthProvider(),
                     ssl_options=SSLOptions(),
                 )
-                thrift_backend._results_message_to_execute_response = Mock(
-                    return_value=(Mock(), Mock())
-                )
+                thrift_backend._results_message_to_execute_response = Mock()
 
                 thrift_backend._handle_execute_response(execute_resp, Mock())
 
@@ -957,12 +948,6 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
             operationHandle=self.operation_handle,
         )
 
-        # Mock the operation status response
-        op_state = ttypes.TGetOperationStatusResp(
-            status=self.okay_status,
-            operationState=ttypes.TOperationState.FINISHED_STATE,
-        )
-        tcli_service_instance.GetOperationStatus.return_value = op_state
         tcli_service_instance.GetResultSetMetadata.return_value = (
             t_get_result_set_metadata_resp
         )
@@ -977,7 +962,7 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
             t_execute_resp, Mock()
         )
 
-        self.assertEqual(arrow_schema_bytes, arrow_schema_mock)
+        self.assertEqual(execute_response.arrow_schema_bytes, arrow_schema_mock)
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     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):
             operationHandle=self.operation_handle,
         )
 
-        # Mock the operation status response
-        op_state = ttypes.TGetOperationStatusResp(
-            status=self.okay_status,
-            operationState=ttypes.TOperationState.FINISHED_STATE,
-        )
-        tcli_service_instance.GetOperationStatus.return_value = op_state
         tcli_service_instance.GetResultSetMetadata.return_value = hive_schema_req
         tcli_service_instance.GetOperationStatus.return_value = (
             ttypes.TGetOperationStatusResp(
@@ -1694,9 +1673,7 @@ def test_handle_execute_response_sets_active_op_handle(self):
         thrift_backend = self._make_fake_thrift_backend()
         thrift_backend._check_direct_results_for_error = Mock()
         thrift_backend._wait_until_command_done = Mock()
-        thrift_backend._results_message_to_execute_response = Mock(
-            return_value=(Mock(), Mock())
-        )
+        thrift_backend._results_message_to_execute_response = Mock()
 
         # Create a mock response with a real operation handle
         mock_resp = Mock()
@@ -2256,8 +2233,7 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
     @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     @patch(
-        "databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response",
-        return_value=(Mock(), Mock()),
+        "databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response"
     )
     def test_execute_command_sets_complex_type_fields_correctly(
         self, mock_handle_execute_response, tcli_service_class, mock_result_set