@@ -113,10 +113,12 @@ def connection(self, extra_params=()):
113113conn .close ()
114114
115115@contextmanager
116- def cursor (self ,extra_params = ()):
116+ def cursor (self ,extra_params = (), extra_cursor_params = () ):
117117with self .connection (extra_params )as conn :
118118cursor = conn .cursor (
119- arraysize = self .arraysize ,buffer_size_bytes = self .buffer_size_bytes
119+ arraysize = self .arraysize ,
120+ buffer_size_bytes = self .buffer_size_bytes ,
121+ ** extra_cursor_params ,
120122 )
121123try :
122124yield cursor
@@ -945,6 +947,60 @@ def test_result_set_close(self):
945947finally :
946948cursor .close ()
947949
950+ def test_row_limit_with_larger_result (self ):
951+ """Test that row_limit properly constrains results when query would return more rows"""
952+ row_limit = 1000
953+ with self .cursor (extra_cursor_params = {"row_limit" :row_limit })as cursor :
954+ # Execute a query that returns more than row_limit rows
955+ cursor .execute ("SELECT * FROM range(2000)" )
956+ rows = cursor .fetchall ()
957+
958+ # Check if the number of rows is limited to row_limit
959+ assert len (rows )== row_limit ,f"Expected{ row_limit } { len (rows )}
960+
961+ def test_row_limit_with_smaller_result (self ):
962+ """Test that row_limit doesn't affect results when query returns fewer rows than limit"""
963+ row_limit = 100
964+ expected_rows = 50
965+ with self .cursor (extra_cursor_params = {"row_limit" :row_limit })as cursor :
966+ # Execute a query that returns fewer than row_limit rows
967+ cursor .execute (f"SELECT * FROM range({ expected_rows } )
968+ rows = cursor .fetchall ()
969+
970+ # Check if all rows are returned (not limited by row_limit)
971+ assert (
972+ len (rows )== expected_rows
973+ ),f"Expected{ expected_rows } { len (rows )}
974+
975+ @skipUnless (pysql_supports_arrow (),"arrow test needs arrow support" )
976+ def test_row_limit_with_arrow_larger_result (self ):
977+ """Test that row_limit properly constrains arrow results when query would return more rows"""
978+ row_limit = 800
979+ with self .cursor (extra_cursor_params = {"row_limit" :row_limit })as cursor :
980+ # Execute a query that returns more than row_limit rows
981+ cursor .execute ("SELECT * FROM range(1500)" )
982+ arrow_table = cursor .fetchall_arrow ()
983+
984+ # Check if the number of rows in the arrow table is limited to row_limit
985+ assert (
986+ arrow_table .num_rows == row_limit
987+ ),f"Expected{ row_limit } { arrow_table .num_rows }
988+
989+ @skipUnless (pysql_supports_arrow (),"arrow test needs arrow support" )
990+ def test_row_limit_with_arrow_smaller_result (self ):
991+ """Test that row_limit doesn't affect arrow results when query returns fewer rows than limit"""
992+ row_limit = 200
993+ expected_rows = 100
994+ with self .cursor (extra_cursor_params = {"row_limit" :row_limit })as cursor :
995+ # Execute a query that returns fewer than row_limit rows
996+ cursor .execute (f"SELECT * FROM range({ expected_rows } )
997+ arrow_table = cursor .fetchall_arrow ()
998+
999+ # Check if all rows are returned (not limited by row_limit)
1000+ assert (
1001+ arrow_table .num_rows == expected_rows
1002+ ),f"Expected{ expected_rows } { arrow_table .num_rows }
1003+
9481004
9491005# use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep
9501006# the 429/503 subsuites separate since they execute under different circumstances.