Merge branch 'exec-resp-norm' into cloudfetch-sea

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -756,11 +756,13 @@ def _results_message_to_execute_response(self, resp, operation_state):
             )
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
+
         has_more_rows = (
             (not direct_results)
             or (not direct_results.resultSet)
             or direct_results.resultSet.hasMoreRows
         )
+
         description = self._hive_schema_to_description(
             t_result_set_metadata_resp.schema
         )
@@ -782,18 +784,19 @@ def _results_message_to_execute_response(self, resp, operation_state):
         if status is None:
             raise ValueError(f"Unknown command state: {operation_state}")
 
-        return ExecuteResponse(
+        execute_response = ExecuteResponse(
             command_id=command_id,
             status=status,
             description=description,
-            has_more_rows=has_more_rows,
             has_been_closed_server_side=has_been_closed_server_side,
             lz4_compressed=lz4_compressed,
             is_staging_operation=t_result_set_metadata_resp.isStagingOperation,
             arrow_schema_bytes=schema_bytes,
             result_format=t_result_set_metadata_resp.resultFormat,
         )
 
+        return execute_response, has_more_rows
+
     def get_execution_result(
         self, command_id: CommandId, cursor: "Cursor"
     ) -> "ResultSet":
@@ -842,7 +845,6 @@ def get_execution_result(
             command_id=command_id,
             status=status,
             description=description,
-            has_more_rows=has_more_rows,
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=is_staging_operation,
@@ -860,6 +862,7 @@ def get_execution_result(
             t_row_set=resp.results,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -972,7 +975,9 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response = self._handle_execute_response(resp, cursor)
+            execute_response, has_more_rows = self._handle_execute_response(
+                resp, cursor
+            )
 
             t_row_set = None
             if resp.directResults and resp.directResults.resultSet:
@@ -988,6 +993,7 @@ def execute_command(
                 t_row_set=t_row_set,
                 max_download_threads=self.max_download_threads,
                 ssl_options=self._ssl_options,
+                has_more_rows=has_more_rows,
             )
 
     def get_catalogs(
@@ -1009,7 +1015,7 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1025,6 +1031,7 @@ def get_catalogs(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def get_schemas(
@@ -1050,7 +1057,7 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1066,6 +1073,7 @@ def get_schemas(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def get_tables(
@@ -1095,7 +1103,7 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1111,6 +1119,7 @@ def get_tables(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def get_columns(
@@ -1140,7 +1149,7 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, has_more_rows = self._handle_execute_response(resp, cursor)
 
         t_row_set = None
         if resp.directResults and resp.directResults.resultSet:
@@ -1156,6 +1165,7 @@ def get_columns(
             t_row_set=t_row_set,
             max_download_threads=self.max_download_threads,
             ssl_options=self._ssl_options,
+            has_more_rows=has_more_rows,
         )
 
     def _handle_execute_response(self, resp, cursor):

src/databricks/sql/backend/types.py

@@ -402,7 +402,6 @@ class ExecuteResponse:
     command_id: CommandId
     status: CommandState
     description: Optional[List[Tuple]] = None
-    has_more_rows: bool = False
     has_been_closed_server_side: bool = False
     lz4_compressed: bool = True
     is_staging_operation: bool = False

src/databricks/sql/result_set.py

@@ -163,6 +163,7 @@ def __init__(
         t_row_set=None,
         max_download_threads: int = 10,
         ssl_options=None,
+        has_more_rows: bool = True,
     ):
         """
         Initialize a ThriftResultSet with direct access to the ThriftDatabricksClient.
@@ -177,6 +178,7 @@ def __init__(
             t_row_set: The TRowSet containing result data (if available)
             max_download_threads: Maximum number of download threads for cloud fetch
             ssl_options: SSL options for cloud fetch
+            has_more_rows: Whether there are more rows to fetch
         """
         # Initialize ThriftResultSet-specific attributes
         self._arrow_schema_bytes = execute_response.arrow_schema_bytes
@@ -208,7 +210,7 @@ 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,
+            has_more_rows=has_more_rows,
             results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,

tests/unit/test_client.py

@@ -213,7 +213,6 @@ def test_closing_result_set_hard_closes_commands(self):
         type(mock_connection).session = PropertyMock(return_value=mock_session)
 
         mock_thrift_backend.fetch_results.return_value = (Mock(), False)
-
         result_set = ThriftResultSet(
             mock_connection, mock_results_response, mock_thrift_backend
         )
@@ -479,7 +478,6 @@ def make_fake_row_slice(n_rows):
         mock_aq = Mock()
         mock_aq.next_n_rows.side_effect = make_fake_row_slice
         mock_thrift_backend.execute_command.return_value.arrow_queue = mock_aq
-        mock_thrift_backend.fetch_results.return_value = (mock_aq, True)
 
         cursor = client.Cursor(Mock(), mock_thrift_backend)
         cursor.execute("foo")

tests/unit/test_fetches.py

@@ -57,17 +57,13 @@ def make_dummy_result_set_from_initial_results(initial_results):
                 command_id=None,
                 status=None,
                 has_been_closed_server_side=True,
-                has_more_rows=False,
                 description=description,
                 lz4_compressed=True,
                 is_staging_operation=False,
             ),
             thrift_client=mock_thrift_backend,
             t_row_set=None,
         )
-
-        # Replace the results queue with our arrow_queue
-        rs.results = arrow_queue
         return rs
 
     @staticmethod
@@ -105,7 +101,6 @@ def fetch_results(
                 command_id=None,
                 status=None,
                 has_been_closed_server_side=False,
-                has_more_rows=True,
                 description=description,
                 lz4_compressed=True,
                 is_staging_operation=False,