return has_more_rows from ExecResponse conversion during GetRespMetadata

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -758,6 +758,12 @@ 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
         )
@@ -779,7 +785,7 @@ 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,
@@ -790,6 +796,8 @@ def _results_message_to_execute_response(self, resp, operation_state):
             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":
@@ -855,6 +863,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):
@@ -967,7 +976,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:
@@ -983,6 +994,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(
@@ -1004,7 +1016,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:
@@ -1020,6 +1032,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(
@@ -1045,7 +1058,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:
@@ -1061,6 +1074,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(
@@ -1090,7 +1104,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:
@@ -1106,6 +1120,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(
@@ -1135,7 +1150,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:
@@ -1151,6 +1166,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/result_set.py

@@ -160,6 +160,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.
@@ -174,6 +175,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
@@ -205,7 +207,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=True,
+            has_more_rows=has_more_rows,
             results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,