pass CommandId to ExecResp

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -352,6 +352,7 @@ def make_request(self, method, request, retryable=True):
         Will stop retry attempts if total elapsed time + next retry delay would exceed
         _retry_stop_after_attempts_duration.
         """
+
         # basic strategy: build range iterator rep'ing number of available
         # retries. bounds can be computed from there. iterate over it with
         # retries until success or final failure achieved.
@@ -866,9 +867,13 @@ def get_execution_result(
             ssl_options=self._ssl_options,
         )
 
+        status = CommandState.from_thrift_state(resp.status)
+        if status is None:
+            raise ValueError(f"Invalid operation state: {resp.status}")
+
         execute_response = ExecuteResponse(
             command_id=command_id,
-            status=resp.status,
+            status=status,
             description=description,
             has_more_rows=has_more_rows,
             results_queue=queue,

src/databricks/sql/result_set.py

@@ -207,7 +207,7 @@ def _fill_results_buffer(self):
             use_cloud_fetch=self._use_cloud_fetch,
         )
         self.results = results
-        self._has_more_rows = has_more_rows
+        self.has_more_rows = has_more_rows
 
     def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]
@@ -291,7 +291,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         while (
             n_remaining_rows > 0
             and not self.has_been_closed_server_side
-            and self._has_more_rows
+            and self.has_more_rows
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
@@ -316,7 +316,7 @@ def fetchmany_columnar(self, size: int):
         while (
             n_remaining_rows > 0
             and not self.has_been_closed_server_side
-            and self._has_more_rows
+            and self.has_more_rows
         ):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
@@ -331,7 +331,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
 
-        while not self.has_been_closed_server_side and self._has_more_rows:
+        while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             if isinstance(results, ColumnTable) and isinstance(
@@ -357,7 +357,7 @@ def fetchall_columnar(self):
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
 
-        while not self.has_been_closed_server_side and self._has_more_rows:
+        while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             results = self.merge_columnar(results, partial_results)