fix: harden async result rowcount retrieval

Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>

Author: aakash-saravanan-db

src/databricks/sql/backend/thrift_backend.py

@@ -927,6 +927,11 @@ def _results_message_to_execute_response(
     def get_execution_result(
         self, command_id: CommandId, cursor: Cursor
     ) -> "ResultSet":
+        """Fetch a completed async result and its latest operation metadata.
+
+        FetchResults does not carry ``numModifiedRows``, so Thrift refreshes
+        operation status before constructing the result set.
+        """
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift command ID")

src/databricks/sql/client.py

@@ -1701,7 +1701,9 @@ def get_async_execution_result(self):
                 self.active_command_id, self
             )
 
-            num_modified_rows = self.active_result_set.num_modified_rows
+            num_modified_rows = getattr(
+                self.active_result_set, "num_modified_rows", None
+            )
             self.rowcount = num_modified_rows if num_modified_rows is not None else -1
 
             if self.active_result_set and self.active_result_set.is_staging_operation:

tests/unit/test_thrift_backend.py

@@ -179,6 +179,44 @@ def test_async_result_reads_status_count_and_preserves_rows(count):
         assert result_set.call_args.kwargs["t_row_set"] is sentinel.counter_rows
 
 
+@pytest.mark.parametrize(
+    "operation_state,exception_type",
+    [
+        (ttypes.TOperationState.ERROR_STATE, ServerOperationError),
+        (ttypes.TOperationState.CLOSED_STATE, DatabaseError),
+    ],
+)
+def test_async_result_rejects_terminal_error_state_before_fetch(
+    operation_state, exception_type
+):
+    with patch(
+        "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True
+    ):
+        backend = ThriftDatabricksClient(
+            "foobar",
+            443,
+            "path",
+            [],
+            auth_provider=AuthProvider(),
+            ssl_options=SSLOptions(),
+            http_client=MagicMock(),
+        )
+        backend._poll_for_status = Mock(
+            return_value=ttypes.TGetOperationStatusResp(
+                status=ThriftBackendTestSuite.okay_status,
+                operationState=operation_state,
+                errorMessage="async operation failed",
+            )
+        )
+        backend.make_request = Mock()
+        cid = CommandId.from_thrift_handle(ThriftBackendTestSuite.operation_handle)
+
+        with pytest.raises(exception_type):
+            backend.get_execution_result(cid, Mock())
+
+        backend.make_request.assert_not_called()
+
+
 @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
 class ThriftBackendTestSuite(unittest.TestCase):
     okay_status = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS)