test: focus async rowcount regression coverage

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

Author: aakash-saravanan-db

KERNEL_REV

@@ -1 +1 @@
-426993a6838aebb7d37b8fc89fa606b4897dd3f2
+453017f8c305e083e14b9b85183bc00920523ad5

tests/unit/test_kernel_client.py

@@ -1234,53 +1234,21 @@ def test_get_query_state_uses_reader_after_result_stream_started():
 
 
 @pytest.mark.parametrize("count", [7, 0, None])
-@pytest.mark.parametrize("source", ["owner", "reader", "attached"])
-def test_async_result_count_after_await(source, count):
+def test_get_execution_result_uses_retained_owning_handle_first(count):
+    """The first in-process result fetch uses the retained submitting
+    handle so the kernel finalizes the original async statement telemetry."""
     c = _make_client()
     c._kernel_session = MagicMock()
-    cid = CommandId.from_sea_statement_id("async-count")
-    owner = MagicMock()
-    handle = owner
-    if source == "attached":
-        c._kernel_session.attach_async_statement.return_value = handle
-    else:
-        c._async_handles[cid.guid] = owner
-        if source == "reader":
-            c._async_result_stream_started.add(cid.guid)
-            handle = owner.reader.return_value
-
+    fake_stream = MagicMock()
+    fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
+    handle = MagicMock()
     del handle.num_modified_rows
-    stream = MagicMock()
-    stream.arrow_schema.return_value = pa.schema([("num_affected_rows", pa.int64())])
-    stream.fetch_next_batch.side_effect = [
-        pa.record_batch(
-            [[count if count is not None else 1]], names=["num_affected_rows"]
-        ),
-        None,
-    ]
 
     def await_result():
         handle.num_modified_rows = count
-        return stream
+        return fake_stream
 
     handle.await_result.side_effect = await_result
-    cursor = MagicMock(arraysize=100, buffer_size_bytes=1024)
-    rs = c.get_execution_result(cid, cursor)
-
-    assert rs.num_modified_rows == count
-    stream.fetch_next_batch.assert_not_called()
-    assert [list(row) for row in rs.fetchall()] == [[count if count is not None else 1]]
-
-
-def test_get_execution_result_uses_retained_owning_handle_first():
-    """The first in-process result fetch uses the retained submitting
-    handle so the kernel finalizes the original async statement telemetry."""
-    c = _make_client()
-    c._kernel_session = MagicMock()
-    fake_stream = MagicMock()
-    fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
-    handle = MagicMock()
-    handle.await_result.return_value = fake_stream
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1291,6 +1259,8 @@ def test_get_execution_result_uses_retained_owning_handle_first():
     rs = c.get_execution_result(cid, cursor=cursor)
 
     assert rs is not None
+    assert rs.num_modified_rows == count
+    fake_stream.fetch_next_batch.assert_not_called()
     c._kernel_session.attach_async_statement.assert_not_called()
     handle.await_result.assert_called_once_with()
     assert cid.guid in c._async_result_stream_started
@@ -1303,6 +1273,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     handle = _attach_returns(c, await_result=fake_stream)
+    handle.num_modified_rows = 7
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1311,6 +1282,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle():
     rs = c.get_execution_result(cid, cursor=cursor)
 
     assert rs is not None
+    assert rs.num_modified_rows == 7
     c._kernel_session.attach_async_statement.assert_called_once_with("async-1")
     handle.await_result.assert_called_once_with()
 
@@ -1365,6 +1337,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
     retry_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
     reader = MagicMock()
     reader.await_result.return_value = retry_stream
+    reader.num_modified_rows = 0
     owning_handle.reader.return_value = reader
     cursor = MagicMock()
     cursor.arraysize = 100
@@ -1373,6 +1346,7 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade
     rs = c.get_execution_result(cid, cursor=cursor)
 
     assert rs is not None
+    assert rs.num_modified_rows == 0
     owning_handle.reader.assert_called_once_with()
     c._kernel_session.attach_async_statement.assert_not_called()
     reader.await_result.assert_called_once_with()

tests/unit/test_thrift_backend.py

@@ -134,113 +134,6 @@ def test_metadata_catalog_escaping(
             assert getattr(request, thrift_field) == value
 
 
-@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
-@pytest.mark.parametrize(
-    "operation_state,count,expected_count",
-    [
-        (ttypes.TOperationState.FINISHED_STATE, 7, 7),
-        (ttypes.TOperationState.FINISHED_STATE, 0, 0),
-        (ttypes.TOperationState.FINISHED_STATE, None, None),
-        (ttypes.TOperationState.RUNNING_STATE, 7, None),
-    ],
-)
-@pytest.mark.parametrize("cache_status", [True, False])
-def test_async_result_reads_only_terminal_status_count_and_preserves_rows(
-    operation_state, count, expected_count, cache_status
-):
-    with (
-        patch(
-            "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True
-        ),
-        patch("databricks.sql.backend.thrift_backend.ThriftResultSet") as result_set,
-    ):
-        backend = ThriftDatabricksClient(
-            "foobar",
-            443,
-            "path",
-            [],
-            auth_provider=AuthProvider(),
-            ssl_options=SSLOptions(),
-            http_client=MagicMock(),
-        )
-        operation_status = ttypes.TGetOperationStatusResp(
-            status=ThriftBackendTestSuite.okay_status,
-            operationState=operation_state,
-            numModifiedRows=count,
-        )
-        backend._poll_for_status = Mock(return_value=operation_status)
-        metadata = ttypes.TGetResultSetMetadataResp(
-            schema=ttypes.TTableSchema(columns=[]),
-            arrowSchema=pyarrow.schema([]).serialize().to_pybytes(),
-            resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET,
-        )
-        backend.make_request = Mock(
-            return_value=Mock(
-                resultSetMetadata=metadata,
-                results=sentinel.counter_rows,
-                hasMoreRows=False,
-            )
-        )
-        cid = CommandId.from_thrift_handle(ThriftBackendTestSuite.operation_handle)
-        if cache_status:
-            expected_state = (
-                CommandState.SUCCEEDED
-                if operation_state == ttypes.TOperationState.FINISHED_STATE
-                else CommandState.RUNNING
-            )
-            assert backend.get_query_state(cid) == expected_state
-        backend.get_execution_result(cid, Mock())
-
-        expected_polls = (
-            2
-            if cache_status and operation_state == ttypes.TOperationState.RUNNING_STATE
-            else 1
-        )
-        assert backend._poll_for_status.call_count == expected_polls
-        backend._poll_for_status.assert_called_with(cid.to_thrift_handle())
-        response = result_set.call_args.kwargs["execute_response"]
-        assert response.num_modified_rows == expected_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)
@@ -320,6 +213,64 @@ def _make_fake_thrift_backend(self):
         thrift_backend._create_arrow_table.return_value = (MagicMock(), Mock())
         return thrift_backend
 
+    @patch("databricks.sql.backend.thrift_backend.ThriftResultSet")
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
+    def test_async_result_rowcount(self, tcli_service_class, result_set):
+        for state, count, expected_count, polled in [
+            (ttypes.TOperationState.FINISHED_STATE, 7, 7, True),
+            (ttypes.TOperationState.FINISHED_STATE, 0, 0, True),
+            (ttypes.TOperationState.FINISHED_STATE, None, None, True),
+            (ttypes.TOperationState.FINISHED_STATE, 7, 7, False),
+            (ttypes.TOperationState.RUNNING_STATE, 7, None, False),
+        ]:
+            with self.subTest(state=state, count=count, polled=polled):
+                backend = self._make_fake_thrift_backend()
+                backend._poll_for_status = Mock(
+                    return_value=ttypes.TGetOperationStatusResp(
+                        status=self.okay_status,
+                        operationState=state,
+                        numModifiedRows=count,
+                    )
+                )
+                backend.make_request = Mock(
+                    return_value=Mock(
+                        resultSetMetadata=self.metadata_resp,
+                        results=sentinel.counter_rows,
+                        hasMoreRows=False,
+                    )
+                )
+                cid = CommandId.from_thrift_handle(self.operation_handle)
+                if polled:
+                    self.assertEqual(
+                        backend.get_query_state(cid), CommandState.SUCCEEDED
+                    )
+                backend.get_execution_result(cid, Mock())
+
+                backend._poll_for_status.assert_called_once_with(cid.to_thrift_handle())
+                response = result_set.call_args.kwargs["execute_response"]
+                self.assertEqual(response.num_modified_rows, expected_count)
+                self.assertIs(
+                    result_set.call_args.kwargs["t_row_set"], sentinel.counter_rows
+                )
+
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
+    def test_async_result_checks_status_before_fetch(self, tcli_service_class):
+        backend = self._make_fake_thrift_backend()
+        backend._poll_for_status = Mock(
+            return_value=ttypes.TGetOperationStatusResp(
+                status=self.okay_status,
+                operationState=ttypes.TOperationState.ERROR_STATE,
+                errorMessage="async operation failed",
+            )
+        )
+        backend.make_request = Mock()
+        cid = CommandId.from_thrift_handle(self.operation_handle)
+
+        with self.assertRaises(ServerOperationError):
+            backend.get_execution_result(cid, Mock())
+
+        backend.make_request.assert_not_called()
+
     def test_hive_schema_to_arrow_schema_preserves_column_names(self):
         columns = [
             ttypes.TColumnDesc(