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