ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4419356673 at src/databricks/sql/result_set.py:382
  - #4419356679 at src/databricks/sql/backend/sea/result_set.py:178

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

Author: peco-engineer-bot[bot]

src/databricks/sql/backend/sea/utils/filters.py

@@ -174,8 +174,9 @@ def _filter_arrow_result_set(
             raise ValueError(f"Column index {column_index} is out of bounds")
         column_name = result_set.description[column_index][0]
 
-        # Get all remaining rows as Arrow table and filter it
-        arrow_table = result_set.results.remaining_rows()
+        # Get all remaining rows as Arrow table and filter it. Go through
+        # fetchall_arrow() so any rows buffered by fetchone() are included.
+        arrow_table = result_set.fetchall_arrow()
         filtered_table = ResultSetFilter._filter_arrow_table(
             arrow_table, column_name, allowed_values, case_sensitive
         )
@@ -220,8 +221,9 @@ def _filter_json_result_set(
         if column_index >= len(result_set.description):
             raise ValueError(f"Column index {column_index} is out of bounds")
 
-        # Extract rows
-        all_rows = result_set.results.remaining_rows()
+        # Extract rows. Go through fetchall_json() so any rows buffered by
+        # fetchone() are included.
+        all_rows = result_set.fetchall_json()
 
         # Convert allowed values if case-insensitive
         if not case_sensitive:

src/databricks/sql/result_set.py

@@ -373,10 +373,7 @@ def _fetchone_buffered(
         """
         if self._row_buffer_count() == 0:
             cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
-            target = getattr(
-                self, "_fetchone_refill_rows", _FETCHONE_BUFFER_INITIAL_ROWS
-            )
-            size = max(min(target, cap), 1)
+            size = max(min(self._fetchone_refill_rows, cap), 1)
             self._fetchone_refill_rows = min(
                 size * _FETCHONE_BUFFER_GROWTH_FACTOR, _FETCHONE_BUFFER_MAX_ROWS
             )
@@ -396,12 +393,15 @@ def _fetchone_buffered(
                 # Warn once per result set; later failing batches (e.g. a column
                 # type that fails in every row) log at DEBUG to avoid a WARNING
                 # per refill for callers that catch and keep fetching.
-                warned = getattr(self, "_fetchone_conversion_warned", False)
                 fallback = self._convert_rows_individually(
                     raw,
                     convert,
                     batch_error,
-                    log_level=logging.DEBUG if warned else logging.WARNING,
+                    log_level=(
+                        logging.DEBUG
+                        if self._fetchone_conversion_warned
+                        else logging.WARNING
+                    ),
                 )
                 if fallback is None:
                     raise

tests/unit/test_filters.py

@@ -16,9 +16,9 @@ def setUp(self):
         # Create a mock SeaResultSet
         self.mock_sea_result_set = MagicMock()
 
-        # Set up the remaining_rows method on the results attribute
+        # Set up fetchall_json, which the filters use to read all remaining rows
         self.mock_sea_result_set.results = MagicMock()
-        self.mock_sea_result_set.results.remaining_rows.return_value = [
+        self.mock_sea_result_set.fetchall_json.return_value = [
             ["catalog1", "schema1", "table1", "owner1", "2023-01-01", "TABLE", ""],
             ["catalog1", "schema1", "table2", "owner1", "2023-01-01", "VIEW", ""],
             [