fix: fetchmany supports irregular sizes

allow size of fetchmany to vary per call

fixes a bug in jupysql where only the first 2 rows are fetched

fakesnow/fakes.py

@@ -391,12 +391,13 @@ def fetchmany(self, size: int | None = None) -> list[tuple] | list[dict]:
         if self._arrow_table is None:
             # mimic snowflake python connector error type
             raise TypeError("No open result set")
+        tslice = self._arrow_table.slice(offset=self._arrow_table_fetch_index or 0, length=size).to_pylist()
+
         if self._arrow_table_fetch_index is None:
-            self._arrow_table_fetch_index = 0
+            self._arrow_table_fetch_index = size
         else:
             self._arrow_table_fetch_index += size
 
-        tslice = self._arrow_table.slice(offset=self._arrow_table_fetch_index, length=size).to_pylist()
         return tslice if self._use_dict_result else [tuple(d.values()) for d in tslice]
 
     def get_result_batches(self) -> list[ResultBatch] | None:

tests/test_fakes.py

@@ -651,20 +651,21 @@ def test_fetchmany(conn: snowflake.connector.SnowflakeConnection):
         cur.execute("insert into customers values (3, 'Jeremy', 'K')")
         cur.execute("select id, first_name, last_name from customers")
 
+        # mimic jupysql fetchmany behaviour
         assert cur.fetchmany(2) == [(1, "Jenny", "P"), (2, "Jasper", "M")]
-        assert cur.fetchmany(2) == [(3, "Jeremy", "K")]
-        assert cur.fetchmany(2) == []
+        assert cur.fetchmany(5) == [(3, "Jeremy", "K")]
+        assert cur.fetchmany(5) == []
 
     with conn.cursor(snowflake.connector.cursor.DictCursor) as cur:
         cur.execute("select id, first_name, last_name from customers")
         assert cur.fetchmany(2) == [
             {"ID": 1, "FIRST_NAME": "Jenny", "LAST_NAME": "P"},
             {"ID": 2, "FIRST_NAME": "Jasper", "LAST_NAME": "M"},
         ]
-        assert cur.fetchmany(2) == [
+        assert cur.fetchmany(5) == [
             {"ID": 3, "FIRST_NAME": "Jeremy", "LAST_NAME": "K"},
         ]
-        assert cur.fetchmany(2) == []
+        assert cur.fetchmany(5) == []
 
 
 def test_fetch_pandas_all(cur: snowflake.connector.cursor.SnowflakeCursor):