Fix fetchall returning already returned rows

If fetchall was called after fetching some rows already via either
fetchmany or fetchone then some rows were duplicated.

    cur = conn.cursor()
    cur.execute("SELECT * FROM ( VALUES (1), (2), (3), (4), (5), (6))")
    print(cur.fetchmany(2))
    print(cur.fetchall()) # should return 4 rows but returns all rows
    print(cur.fetchmany(10)) # should return no rows but returns next 4 rows

tests/integration/test_dbapi_integration.py

@@ -1328,6 +1328,15 @@ def test_select_tpch_1000(trino_connection):
     assert len(rows) == 1000
 
 
+def test_fetch_cursor(trino_connection):
+    cur = trino_connection.cursor()
+    cur.execute("SELECT * FROM tpch.sf1.customer LIMIT 1000")
+    for _ in range(100):
+        cur.fetchone()
+    assert len(cur.fetchmany(400)) == 400
+    assert len(cur.fetchall()) == 500
+
+
 def test_cancel_query(trino_connection):
     cur = trino_connection.cursor()
     cur.execute("SELECT * FROM tpch.sf1.customer")

trino/dbapi.py

@@ -23,6 +23,7 @@
 import uuid
 from collections import OrderedDict
 from decimal import Decimal
+from itertools import islice
 from threading import Lock
 from time import time
 from typing import Any, Dict, List, NamedTuple, Optional  # NOQA for mypy types
@@ -661,14 +662,7 @@ def fetchmany(self, size=None) -> List[List[Any]]:
         if size is None:
             size = self.arraysize
 
-        result = []
-        for _ in range(size):
-            row = self.fetchone()
-            if row is None:
-                break
-            result.append(row)
-
-        return result
+        return list(islice(iter(self.fetchone, None), size))
 
     def describe(self, sql: str) -> List[DescribeOutput]:
         """
@@ -696,7 +690,7 @@ def genall(self):
         return self._query.result
 
     def fetchall(self) -> List[List[Any]]:
-        return list(self.genall())
+        return list(iter(self.fetchone, None))
 
     def cancel(self):
         if self._query is None: