@@ -725,19 +725,66 @@ def test_fetchone_convert_failure_warns_once_per_result_set(self):
725725 "_convert_columnar_table" ,
726726 side_effect = ValueError ("bad value" ),
727727 ), self .assertLogs ("databricks.sql.result_set" , level = "DEBUG" ) as logs :
728- # Three refill batches, every row failing the same way.
728+ # One 4-row refill, then (ramp reset after an all-failing batch) eight
729+ # 1-row refills, every row failing the same way.
729730 for _ in range (12 ):
730731 with self .assertRaises (ValueError ):
731732 rs .fetchone ()
732733 summaries = [
733734 r for r in logs .records if "failed to convert" in r .getMessage ()
734735 ]
735- self .assertEqual (len (summaries ), 3 )
736+ self .assertEqual (len (summaries ), 9 )
736737 self .assertEqual (
737738 [r .levelno for r in summaries ],
738- [logging .WARNING , logging . DEBUG , logging .DEBUG ],
739+ [logging .WARNING ] + [ logging .DEBUG ] * 8 ,
739740 )
740741
742+ def test_fetchone_all_failing_batch_resets_refill_ramp (self ):
743+ from unittest .mock import patch
744+
745+ calls = []
746+
747+ def convert (self_ , table ):
748+ calls .append (table .num_rows )
749+ raise ValueError ("systemic failure" )
750+
751+ rs = self .make_columnar_result_set (30 , arraysize = 10 )
752+ with patch .object (
753+ ThriftResultSet ,
754+ "_convert_columnar_table" ,
755+ autospec = True ,
756+ side_effect = convert ,
757+ ), self .assertLogs ("databricks.sql.result_set" , level = "DEBUG" ):
758+ for _ in range (15 ):
759+ with self .assertRaises (ValueError ):
760+ rs .fetchone ()
761+ # First full refill is bisected (2 * 10 - 1 conversions); after that every
762+ # refill is a single row costing exactly one conversion.
763+ self .assertEqual (len (calls ), 2 * 10 - 1 + 5 )
764+ self .assertEqual (calls [- 5 :], [1 ] * 5 )
765+ self .assertEqual (rs .rownumber , 15 )
766+
767+ def test_fetchone_memory_error_is_not_bisected (self ):
768+ from unittest .mock import patch
769+
770+ calls = []
771+
772+ def convert (self_ , table ):
773+ calls .append (table .num_rows )
774+ raise MemoryError ()
775+
776+ rs = self .make_columnar_result_set (10 , arraysize = 4 )
777+ with patch .object (
778+ ThriftResultSet ,
779+ "_convert_columnar_table" ,
780+ autospec = True ,
781+ side_effect = convert ,
782+ ):
783+ with self .assertRaises (MemoryError ):
784+ rs .fetchone ()
785+ self .assertEqual (calls , [4 ])
786+ self .assertEqual (rs ._row_buffer_count (), 0 )
787+
741788 def test_fetchone_convert_failure_on_every_row_bisects_to_single_rows (self ):
742789 from unittest .mock import patch
743790
0 commit comments