From ae40212b1bc92d464008711b13163badd9a11492 Mon Sep 17 00:00:00 2001
From: Lee Miller <lee.miller@tutanota.com>
Date: Wed, 11 Sep 2024 02:35:09 +0300
Subject: [PATCH] Set a row factory instead of nesting the iterator, iterate
 through the cursor

---
 minode/sql.py | 37 ++++++++++++++++---------------------
 1 file changed, 16 insertions(+), 21 deletions(-)

diff --git a/minode/sql.py b/minode/sql.py
index 4ccb6d9..755f194 100644
--- a/minode/sql.py
+++ b/minode/sql.py
@@ -39,15 +39,16 @@ class Inventory():
                 cur.execute('VACUUM')
         cur.execute("INSERT INTO status VALUES ('lastvacuumtime', ?)", (now,))
         self._db.commit()
+        self._db.row_factory = self.__object
 
-    def __objects(self, cur):
-        return (
-            structure.Object(
-                expires, obj_type, version, stream, data, offset,
-                tag=tag, vector=vector)
-            for (vector, expires, obj_type, version, stream, tag, data, offset)
-            in cur.fetchall()
-        )
+    @staticmethod
+    def __object(cursor, row):
+        if len(cursor.description) != 8:
+            return row
+        vector, expires, obj_type, version, stream, tag, data, offset = row
+        return structure.Object(
+            expires, obj_type, version, stream, data, offset,
+            tag=tag, vector=vector)
 
     def cleanup(self):
         """Remove expired objects"""
@@ -74,7 +75,7 @@ class Inventory():
         cur = self._db.execute(
             'SELECT * FROM objects WHERE '  # nosec B608
             + ' AND '.join(clauses), params)
-        return self.__objects(cur)
+        return cur
 
     def vectors_to_send(self, stream=None):
         cur = self._db.execute(
@@ -82,7 +83,7 @@ class Inventory():
             ' ORDER BY random()',
             (int(time.time()), stream or shared.stream)
         )
-        return [v for v, in cur.fetchall()]
+        return [v for v, in cur]
 
     def get(self, vector, default=None):
         try:
@@ -92,11 +93,10 @@ class Inventory():
 
     def keys(self):
         cur = self._db.execute('SELECT vector FROM objects')
-        return (v for v, in cur.fetchall())
+        return (v for v, in cur)
 
     def values(self):
-        cur = self._db.execute('SELECT * FROM objects')
-        return self.__objects(cur)
+        return self._db.execute('SELECT * FROM objects')
 
     def popitem(self):
         if not self.rowid:
@@ -114,16 +114,11 @@ class Inventory():
         return cur.fetchone()
 
     def __getitem__(self, vector):
-        cur = self._db.execute(
-            'SELECT * FROM objects WHERE vector = ?', (vector,))
-        item = cur.fetchone()
+        item = self._db.execute(
+            'SELECT * FROM objects WHERE vector = ?', (vector,)).fetchone()
         if item is None:
             raise KeyError(vector)
-        vector, expires, obj_type, version, stream, tag, data, offset = item
-        return structure.Object(
-            expires, obj_type, version, stream, data, offset,
-            tag=tag, vector=vector
-        )
+        return item
 
     def __delitem__(self, vector):
         with shared.objects_lock:  # KeyError