fix tests

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/filters.py

@@ -69,15 +69,15 @@ def _filter_sea_result_set(
             command_id=command_id,
             status=result_set.status,
             description=result_set.description,
-            has_more_rows=result_set.has_more_rows,
             has_been_closed_server_side=result_set.has_been_closed_server_side,
-            lz4_compressed=False,
+            lz4_compressed=result_set.lz4_compressed,
+            arrow_schema_bytes=result_set.arrow_schema_bytes,
             is_staging_operation=False,
         )
 
         # Create a new ResultData object with filtered data
         from databricks.sql.backend.sea.models.base import ResultData
-        
+
         result_data = ResultData(data=filtered_rows, external_links=None)
 
         # Create a new SeaResultSet with the filtered data
@@ -89,7 +89,7 @@ def _filter_sea_result_set(
             arraysize=result_set.arraysize,
             result_data=result_data,
         )
-        
+
         return filtered_result_set
 
     @staticmethod

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

@@ -4,7 +4,9 @@
 import re
 from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING, Set
 
-from databricks.sql.backend.sea.utils.constants import ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP
+from databricks.sql.backend.sea.utils.constants import (
+    ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
+)
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
@@ -72,6 +74,8 @@ def _filter_session_configuration(
         )
 
     return filtered_session_configuration
+
+
 class SeaDatabricksClient(DatabricksClient):
     """
     Statement Execution API (SEA) implementation of the DatabricksClient interface.
@@ -89,7 +93,7 @@ class SeaDatabricksClient(DatabricksClient):
     CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}/cancel"
     CHUNKS_PATH_WITH_ID = STATEMENT_PATH + "/{}/result/chunks"
     CHUNK_PATH_WITH_ID_AND_INDEX = STATEMENT_PATH + "/{}/result/chunks/{}"
-    
+
     def __init__(
         self,
         server_hostname: str,
@@ -502,7 +506,7 @@ def _results_message_to_execute_response(self, sea_response, command_id):
         # Initialize result_data_obj and manifest_obj
         result_data_obj = None
         manifest_obj = None
-        
+
         result_data = sea_response.get("result", {})
         if result_data:
             # Convert external links
@@ -554,7 +558,7 @@ def _results_message_to_execute_response(self, sea_response, command_id):
             arrow_schema_bytes=schema_bytes,
             result_format=manifest_data.get("format"),
         )
-        
+
         return execute_response, result_data_obj, manifest_obj
 
     def execute_command(
@@ -776,9 +780,11 @@ def get_execution_result(
         from databricks.sql.result_set import SeaResultSet
 
         # Convert the response to an ExecuteResponse and extract result data
-        execute_response, result_data, manifest = self._results_message_to_execute_response(
-            response_data, command_id
-        )
+        (
+            execute_response,
+            result_data,
+            manifest,
+        ) = self._results_message_to_execute_response(response_data, command_id)
 
         return SeaResultSet(
             connection=cursor.connection,

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

@@ -40,7 +40,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
             )
 
         state = CommandState.from_sea_state(status_data.get("state", ""))
-        if state is None: 
+        if state is None:
             raise ValueError(f"Invalid state: {status_data.get('state', '')}")
 
         status = StatementStatus(
@@ -124,7 +124,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
             )
 
         state = CommandState.from_sea_state(status_data.get("state", ""))
-        if state is None: 
+        if state is None:
             raise ValueError(f"Invalid state: {status_data.get('state', '')}")
 
         status = StatementStatus(

src/databricks/sql/backend/thrift_backend.py

@@ -1179,12 +1179,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)
-        execute_response.command_id = command_id
-        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

@@ -6,7 +6,11 @@
 import pandas
 
 from databricks.sql.backend.sea.backend import SeaDatabricksClient
-from databricks.sql.backend.sea.models.base import ExternalLink, ResultData, ResultManifest
+from databricks.sql.backend.sea.models.base import (
+    ExternalLink,
+    ResultData,
+    ResultManifest,
+)
 from databricks.sql.cloud_fetch_queue import SeaCloudFetchQueue
 from databricks.sql.utils import SeaResultSetQueueFactory
 
@@ -44,10 +48,11 @@ def __init__(
         command_id: CommandId,
         status: CommandState,
         has_been_closed_server_side: bool = False,
-        has_more_rows: bool = False,
         results_queue=None,
         description=None,
         is_staging_operation: bool = False,
+        lz4_compressed: bool = False,
+        arrow_schema_bytes: bytes = b"",
     ):
         """
         A ResultSet manages the results of a single command.
@@ -75,9 +80,10 @@ def __init__(
         self.command_id = command_id
         self.status = status
         self.has_been_closed_server_side = has_been_closed_server_side
-        self.has_more_rows = has_more_rows
         self.results = results_queue
         self._is_staging_operation = is_staging_operation
+        self.lz4_compressed = lz4_compressed
+        self._arrow_schema_bytes = arrow_schema_bytes
 
     def __iter__(self):
         while True:
@@ -181,9 +187,8 @@ def __init__(
             has_more_rows: Whether there are more rows to fetch
         """
         # Initialize ThriftResultSet-specific attributes
-        self._arrow_schema_bytes = execute_response.arrow_schema_bytes
         self._use_cloud_fetch = use_cloud_fetch
-        self.lz4_compressed = execute_response.lz4_compressed
+        self.has_more_rows = has_more_rows
 
         # Build the results queue if t_row_set is provided
         results_queue = None
@@ -210,10 +215,11 @@ def __init__(
             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=has_more_rows,
             results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
+            lz4_compressed=execute_response.lz4_compressed,
+            arrow_schema_bytes=execute_response.arrow_schema_bytes,
         )
 
         # Initialize results queue if not provided
@@ -442,6 +448,7 @@ def map_col_type(type_):
             for column in table_schema_message.columns
         ]
 
+
 class SeaResultSet(ResultSet):
     """ResultSet implementation for SEA backend."""
 
@@ -473,24 +480,26 @@ def __init__(
             if execute_response.command_id
             else None
         )
-        
+
         # Build the results queue
         results_queue = None
-        
+
         if result_data:
             from typing import cast, List
-            
+
             # Convert description to the expected format
             desc = None
             if execute_response.description:
                 desc = cast(List[Tuple[Any, ...]], execute_response.description)
-            
+
             results_queue = SeaResultSetQueueFactory.build_queue(
                 result_data,
                 manifest,
                 str(self.statement_id),
                 description=desc,
-                schema_bytes=execute_response.arrow_schema_bytes if execute_response.arrow_schema_bytes else None,
+                schema_bytes=execute_response.arrow_schema_bytes
+                if execute_response.arrow_schema_bytes
+                else None,
                 max_download_threads=sea_client.max_download_threads,
                 ssl_options=sea_client.ssl_options,
                 sea_client=sea_client,
@@ -506,9 +515,10 @@ def __init__(
             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,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
+            lz4_compressed=execute_response.lz4_compressed,
+            arrow_schema_bytes=execute_response.arrow_schema_bytes,
         )
 
         # Initialize queue for result data if not provided

src/databricks/sql/utils.py

@@ -177,9 +177,7 @@ def build_queue(
                     "SEA client is required for EXTERNAL_LINKS disposition"
                 )
             if not manifest:
-                raise ValueError(
-                    "Manifest is required for EXTERNAL_LINKS disposition"
-                )
+                raise ValueError("Manifest is required for EXTERNAL_LINKS disposition")
 
             return SeaCloudFetchQueue(
                 initial_links=sea_result_data.external_links,

tests/unit/test_client.py

@@ -566,7 +566,7 @@ def test_cursor_keeps_connection_alive(self, mock_client_class):
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_staging_operation_response_is_handled(
         self,
-        mock_thrift_client_class,
+        mock_client_class,
         mock_handle_staging_operation,
         mock_execute_response,
     ):

tests/unit/test_fetches_bench.py

@@ -36,11 +36,9 @@ def make_dummy_result_set_from_initial_results(arrow_table):
             execute_response=ExecuteResponse(
                 status=None,
                 has_been_closed_server_side=True,
-                has_more_rows=False,
                 description=Mock(),
                 command_id=None,
-                arrow_queue=arrow_queue,
-                arrow_schema=arrow_table.schema,
+                arrow_schema_bytes=arrow_table.schema,
             ),
         )
         rs.description = [

tests/unit/test_parameters.py

@@ -64,13 +64,7 @@ class TestSessionHandleChecks(object):
         ],
     )
     def test_get_protocol_version_fallback_behavior(self, test_input, expected):
-        properties = (
-            {"serverProtocolVersion": test_input.serverProtocolVersion}
-            if test_input.serverProtocolVersion
-            else {}
-        )
-        session_id = SessionId.from_thrift_handle(test_input.sessionHandle, properties)
-        assert Connection.get_protocol_version(session_id) == expected
+        assert Connection.get_protocol_version(test_input) == expected
 
     @pytest.mark.parametrize(
         "test_input,expected",