diff --git a/fakesnow/fakes.py b/fakesnow/fakes.py index 739fb8a..f600a6a 100644 --- a/fakesnow/fakes.py +++ b/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: diff --git a/tests/test_fakes.py b/tests/test_fakes.py index 79e547c..20f8987 100644 --- a/tests/test_fakes.py +++ b/tests/test_fakes.py @@ -651,9 +651,10 @@ 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") @@ -661,10 +662,10 @@ def test_fetchmany(conn: snowflake.connector.SnowflakeConnection): {"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):