fix: normalize kernel rowcount access

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

Author: aakash-saravanan-db

src/databricks/sql/backend/kernel/client.py

@@ -138,6 +138,14 @@ def _async_statement_reader(owner: Any) -> Any:
     return reader()
 
 
+def _read_num_modified_rows(handle: Any) -> Optional[int]:
+    try:
+        modified = getattr(handle, "num_modified_rows", None)
+        return modified() if callable(modified) else modified
+    except Exception:
+        return None
+
+
 def _catalog_or_none(value: Optional[str]) -> Optional[str]:
     """Map supported all-catalog wildcards to the kernel's unset filter."""
     if value is None or value in ("%", "*"):
@@ -691,12 +699,7 @@ def execute_command(
         # default. ``getattr`` guards against an older kernel wheel that
         # predates the pyo3 getter. NB the Thrift backend also hardcodes
         # ``-1`` here, so this makes the kernel path *exceed* Thrift.
-        try:
-            modified = getattr(executed, "num_modified_rows", None)
-            if callable(modified):
-                modified = modified()
-        except Exception:
-            modified = None
+        modified = _read_num_modified_rows(executed)
         if modified is not None:
             cursor.rowcount = modified
         # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
@@ -969,7 +972,7 @@ def get_execution_result(
             result_set = self._make_result_set(stream, cursor, command_id)
             # Older kernel wheels do not expose this value. Keep them compatible
             # while allowing a newer binding to populate async rowcount.
-            result_set.num_modified_rows = getattr(handle, "num_modified_rows", None)
+            result_set.num_modified_rows = _read_num_modified_rows(handle)
             return result_set
         except Exception as exc:
             raise _wrap_kernel_exception("get_execution_result", exc) from exc

tests/unit/test_kernel_client.py

@@ -1233,6 +1233,25 @@ def test_get_query_state_uses_reader_after_result_stream_started():
     reader.status.assert_called_once_with()
 
 
+class _RaisingModifiedRows:
+    @property
+    def num_modified_rows(self):
+        raise RuntimeError("unavailable")
+
+
+@pytest.mark.parametrize(
+    ("handle", "expected"),
+    [
+        (types.SimpleNamespace(num_modified_rows=7), 7),
+        (types.SimpleNamespace(num_modified_rows=lambda: 8), 8),
+        (types.SimpleNamespace(), None),
+        (_RaisingModifiedRows(), None),
+    ],
+)
+def test_read_num_modified_rows_compatibility_shapes(handle, expected):
+    assert kernel_client._read_num_modified_rows(handle) == expected
+
+
 @pytest.mark.parametrize("count", [7, 0, None, "unavailable"])
 @pytest.mark.parametrize("source", ["owner", "reader", "attached"])
 def test_async_result_count_after_await(source, count):