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

Addresses:
  - #4419671350 at src/databricks/sql/result_set.py:567

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

Author: peco-engineer-bot[bot]

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

@@ -273,9 +273,10 @@ def close(self) -> None:
         connection-close would otherwise issue a kernel call into an
         already-disposed session.
         """
-        # Drop rows parked by fetchone() so post-close fetches return
-        # nothing, matching the cleared batch FIFO below.
+        # Drop rows parked by fetchone() (and any lost-rows poison) so
+        # post-close fetches return nothing, matching the cleared batch FIFO below.
         self._row_buffer_clear()
+        self._lost_rows_error = None
         if self._kernel_handle is None:
             return
         if not self.connection.open:

src/databricks/sql/result_set.py

@@ -563,8 +563,10 @@ def close(self) -> None:
             if isinstance(e.args[1], CursorAlreadyClosedError):
                 logger.info("Operation was canceled by a prior request")
         finally:
-            # Drop rows parked by fetchone() so post-close fetches return nothing.
+            # Drop rows parked by fetchone() (and any lost-rows poison) so
+            # post-close fetches return nothing.
             self._row_buffer_clear()
+            self._lost_rows_error = None
             self.has_been_closed_server_side = True
             self.status = CommandState.CLOSED
 

tests/unit/test_fetches.py

@@ -390,6 +390,23 @@ def test_close_drops_fetchone_row_buffer(self):
             self.assertEqual(fetch_after_close(rs), [])
             self.assertIsNone(rs.fetchone())
 
+    def test_close_clears_lost_rows_error(self):
+        # A result set poisoned by a lost fetchone() refill must behave like any
+        # other closed result set: post-close fetches return nothing, not raise.
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        rs = self.make_dummy_result_set_from_initial_results([[1], [2], [3]])
+        # Warm refill pulls every row into the buffer, so the queue is empty.
+        rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
+        self.assertSequenceEqual(rs.fetchone(), [1])
+        rs._lost_rows_error = "lost"
+        with self.assertRaises(InternalError):
+            rs.fetchone()
+        rs.close()
+        self.assertIsNone(rs._lost_rows_error)
+        self.assertIsNone(rs.fetchone())
+        self.assertEqual(rs.fetchall(), [])
+
     def test_fetchone_refill_does_not_fetch_extra_pages(self):
         # A refill must not make extra round-trips to fill ``arraysize`` when
         # server pages are smaller (e.g. byte-limited): it serves what the current

tests/unit/test_kernel_result_set.py

@@ -288,6 +288,16 @@ def test_close_drops_fetchone_row_buffer(int_schema):
     assert rs.fetchall() == []
 
 
+def test_close_clears_lost_rows_error(int_schema):
+    # A poisoned result set must return nothing after close(), not keep raising.
+    rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
+    rs._lost_rows_error = "lost"
+    rs.close()
+    assert rs._lost_rows_error is None
+    assert rs.fetchone() is None
+    assert rs.fetchall() == []
+
+
 def test_empty_stream(int_schema):
     rs = _make_rs(_FakeKernelHandle(int_schema, []))
     assert rs.fetchone() is None