Skip to content

Commit 37a9c53

Browse files
ai: apply changes for #1007 (2 review threads)
Addresses: - #4419518440 at src/databricks/sql/result_set.py:417 - #4419518444 at src/databricks/sql/result_set.py:258 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent a6d6ee5 commit 37a9c53

2 files changed

Lines changed: 59 additions & 2 deletions

File tree

‎src/databricks/sql/result_set.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -381,10 +381,12 @@ def _fetchone_buffered(
381381
if self._row_buffer_count() == 0:
382382
cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
383383
size = max(min(self._fetchone_refill_rows, cap), 1)
384+
raw = fetch_raw(size)
385+
# Grow the ramp only once the fetch succeeded, so a transient fetch
386+
# error (and a retried fetchone()) does not skip a ramp step.
384387
self._fetchone_refill_rows = min(
385388
size * _FETCHONE_BUFFER_GROWTH_FACTOR, _FETCHONE_BUFFER_MAX_ROWS
386389
)
387-
raw = fetch_raw(size)
388390
try:
389391
rows = convert(raw)
390392
except _NON_DATA_CONVERSION_ERRORS:
@@ -436,7 +438,15 @@ def _fetchone_buffered(
436438
# that yields [] for a non-empty chunk raises instead of silently
437439
# ending iteration and losing those rows.
438440
raw_count = self._raw_row_count(raw)
439-
if raw_count is not None and raw_count != len(rows):
441+
if raw_count is None:
442+
# Every real queue type (Arrow Table, ColumnTable, list) has a row
443+
# count, so an unknown chunk type is itself a bug. Fail loudly
444+
# rather than skip the drift check and risk rownumber misalignment.
445+
raise InternalError(
446+
"fetchone(): cannot determine the row count of a raw chunk of "
447+
"type {}".format(type(raw).__name__)
448+
)
449+
if raw_count != len(rows):
440450
# fetch_raw advanced _next_row_index (and the Thrift fetch offset)
441451
# by raw_count; the buffer must line up 1:1 with the raw chunk or
442452
# rownumber and fetchmany_arrow slicing would silently drift.

‎tests/unit/test_fetches.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -869,6 +869,30 @@ def convert(self_, table):
869869
self.assertLessEqual(sum(calls), n * (n - 1).bit_length() + n)
870870
self.assertEqual(rs._row_buffer_count(), n - 1)
871871

872+
def test_fetchone_failed_fetch_does_not_advance_refill_ramp(self):
873+
from unittest.mock import patch
874+
875+
rs = self.make_columnar_result_set(100, arraysize=100, warm=False)
876+
initial = rs._fetchone_refill_rows
877+
real_next_n_rows = rs.results.next_n_rows
878+
with patch.object(
879+
rs.results,
880+
"next_n_rows",
881+
side_effect=[RuntimeError("transient"), real_next_n_rows],
882+
):
883+
with self.assertRaises(RuntimeError):
884+
rs.fetchone()
885+
# The failed fetch must not grow the ramp, so a retry refills the same size.
886+
self.assertEqual(rs._fetchone_refill_rows, initial)
887+
sizes = []
888+
with patch.object(
889+
rs.results,
890+
"next_n_rows",
891+
side_effect=lambda n: sizes.append(n) or real_next_n_rows(n),
892+
):
893+
self.assertEqual(rs.fetchone()[0], 0)
894+
self.assertEqual(sizes, [initial])
895+
872896
def test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows(self):
873897
from unittest.mock import patch
874898

@@ -962,6 +986,29 @@ def test_fetchone_empty_conversion_of_nonempty_chunk_fails_loudly(self):
962986
with self.assertRaises(InternalError):
963987
rs.fetchone()
964988

989+
def test_fetchone_unknown_raw_chunk_type_fails_loudly(self):
990+
from unittest.mock import patch
991+
from databricks.sql.exc import InternalError
992+
993+
class _NoRowCount:
994+
pass
995+
996+
rs = self.make_columnar_result_set(10, arraysize=4)
997+
with patch.object(
998+
ThriftResultSet,
999+
"_fetch_available_from_queue",
1000+
autospec=True,
1001+
return_value=_NoRowCount(),
1002+
), patch.object(
1003+
ThriftResultSet,
1004+
"_convert_columnar_table",
1005+
autospec=True,
1006+
return_value=[("x",)],
1007+
):
1008+
with self.assertRaises(InternalError) as ctx:
1009+
rs.fetchone()
1010+
self.assertIn("_NoRowCount", str(ctx.exception))
1011+
9651012
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
9661013
def test_fetchone_then_fetchall_arrow_keeps_order(self):
9671014
rs = self.make_columnar_result_set(10, arraysize=4)

0 commit comments

Comments
 (0)