Merge "improve fetchmany performance when using deque" into main

edgarrmondragon · Apr 29, 2024 · 7782e2b · 7782e2b
2 parents 6888cf7 + 319304e
commit 7782e2b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 16 deletions.
diff --git a/lib/sqlalchemy/connectors/asyncio.py b/lib/sqlalchemy/connectors/asyncio.py
@@ -11,7 +11,6 @@
 
 import asyncio
 import collections
-import itertools
 import sys
 from typing import Any
 from typing import Deque
@@ -232,11 +231,8 @@ def fetchone(self) -> Optional[Any]:
     def fetchmany(self, size: Optional[int] = None) -> Sequence[Any]:
         if size is None:
             size = self.arraysize
-
-        rr = iter(self._rows)
-        retval = list(itertools.islice(rr, 0, size))
-        self._rows = collections.deque(rr)
-        return retval
+        rr = self._rows
+        return [rr.popleft() for _ in range(min(size, len(rr)))]
 
     def fetchall(self) -> Sequence[Any]:
         retval = list(self._rows)

diff --git a/lib/sqlalchemy/engine/cursor.py b/lib/sqlalchemy/engine/cursor.py
@@ -1252,22 +1252,25 @@ def fetchmany(self, result, dbapi_cursor, size=None):
         if size is None:
             return self.fetchall(result, dbapi_cursor)
 
-        buf = list(self._rowbuffer)
-        lb = len(buf)
+        rb = self._rowbuffer
+        lb = len(rb)
+        close = False
         if size > lb:
             try:
                 new = dbapi_cursor.fetchmany(size - lb)
             except BaseException as e:
                 self.handle_exception(result, dbapi_cursor, e)
             else:
                 if not new:
-                    result._soft_close()
+                    # defer closing since it may clear the row buffer
+                    close = True
                 else:
-                    buf.extend(new)
+                    rb.extend(new)
 
-        result = buf[0:size]
-        self._rowbuffer = collections.deque(buf[size:])
-        return result
+        res = [rb.popleft() for _ in range(min(size, len(rb)))]
+        if close:
+            result._soft_close()
+        return res
 
     def fetchall(self, result, dbapi_cursor):
         try:
@@ -1321,9 +1324,8 @@ def fetchmany(self, result, dbapi_cursor, size=None):
         if size is None:
             return self.fetchall(result, dbapi_cursor)
 
-        buf = list(self._rowbuffer)
-        rows = buf[0:size]
-        self._rowbuffer = collections.deque(buf[size:])
+        rb = self._rowbuffer
+        rows = [rb.popleft() for _ in range(min(size, len(rb)))]
         if not rows:
             result._soft_close()
         return rows