ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4418605049 at src/databricks/sql/result_set.py:224

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>

Author: peco-engineer-bot[bot]

src/databricks/sql/result_set.py

@@ -220,8 +220,24 @@ def _fetchone_buffered(
             # still equals the queue position. Rolling the index back would make it
             # lag the queue and break Thrift's expected_row_start_offset check on
             # the next fetch, and retrying convert() on the same data would just
-            # fail again.
-            rows = convert(raw)
+            # fail again. A warning records how many rows were discarded so a
+            # caller that swallows the error and keeps fetching isn't silently
+            # short of rows.
+            try:
+                rows = convert(raw)
+            except Exception:
+                if raw is None:
+                    dropped = 0
+                elif isinstance(raw, list):
+                    dropped = len(raw)
+                else:
+                    dropped = getattr(raw, "num_rows", 0)
+                logger.warning(
+                    "fetchone(): failed to convert a batch of %d row(s); these rows "
+                    "are discarded and will not be returned by later fetches",
+                    dropped,
+                )
+                raise
             if not rows:
                 return None
             self._row_buffer_raw = raw

tests/unit/test_fetches.py

@@ -534,6 +534,23 @@ def test_mixed_fetches_keep_order_and_rownumber(self):
         self.assertEqual(rs.rownumber, 10)
         self.assertIsNone(rs.fetchone())
 
+    def test_fetchone_convert_failure_logs_dropped_rows_and_reraises(self):
+        from unittest.mock import patch
+
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            side_effect=ValueError("bad value"),
+        ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
+            with self.assertRaises(ValueError):
+                rs.fetchone()
+        self.assertTrue(any("batch of 4 row(s)" in m for m in logs.output))
+        # The failed batch is dropped; the next fetch resumes after it.
+        self.assertEqual(rs._row_buffer_count(), 0)
+        self.assertEqual(rs.rownumber, 4)
+        self.assertEqual(rs.fetchone()[0], 4)
+
     @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
     def test_fetchone_then_fetchall_arrow_keeps_order(self):
         rs = self.make_columnar_result_set(10, arraysize=4)