Merge pull request #14 from AnswerDotAI/cursor_row2dict

Use rowtrace func to convert row to dict

apswutils/db.py

@@ -1,7 +1,7 @@
 # This file is from sqlite-utils and copyright and license is the same as that project
 __all__ = ['Database', 'Queryable', 'Table', 'View']
 
-from .utils import chunks, hash_record, suggest_column_types, types_for_column_types, column_affinity, find_spatialite
+from .utils import chunks, hash_record, suggest_column_types, types_for_column_types, column_affinity, find_spatialite, cursor_row2dict
 from collections import namedtuple
 from collections.abc import Mapping
 from typing import cast, Any, Callable, Dict, Generator, Iterable, Union, Optional, List, Tuple, Iterator
@@ -414,10 +414,8 @@ def query(
           parameters, or a dictionary for ``where id = :id``
         """
         cursor = self.execute(sql, tuple(params or tuple()))
-        try: columns = [c[0] for c in cursor.description]
-        except apsw.ExecutionCompleteError: return []
-        for row in cursor:
-            yield dict(zip(columns, row))
+        cursor.row_trace = cursor_row2dict
+        yield from cursor
 
     def execute(
         self, sql: str, parameters: Optional[Union[Iterable, dict]] = None
@@ -1295,11 +1293,8 @@ def rows_where(
         if offset is not None:
             sql += f" offset {offset}"
         cursor = self.db.execute(sql, where_args or [])
-        # If no records found, return empty list
-        try: columns = [c[0] for c in cursor.description]
-        except apsw.ExecutionCompleteError: return []
-        for row in cursor:
-            yield dict(zip(columns, row))
+        cursor.row_trace = cursor_row2dict
+        yield from cursor
 
     def pks_and_rows_where(
         self,
@@ -2670,10 +2665,8 @@ def search(
             ),
             args,
         )
-        try: columns = [c[0] for c in cursor.description]
-        except apsw.ExecutionCompleteError: return []        
-        for row in cursor:
-            yield dict(zip(columns, row))
+        cursor.row_trace = cursor_row2dict
+        yield from cursor
 
     def value_or_default(self, key, value):
         return self._defaults[key] if value is DEFAULT else value
@@ -2764,11 +2757,8 @@ def update(
         self.result = []
         try:
             cursor = self.db.execute(sql, args)
-            try: columns = [c[0] for c in cursor.description]
-            except apsw.ExecutionCompleteError: return self
-
-            for row in cursor:
-                self.result.append(dict(zip(columns, row)))
+            cursor.row_trace = cursor_row2dict
+            self.result = list(cursor)
         except apsw.SQLError as e:
             if alter and (" column" in e.args[0]):
                 # Attempt to add any missing columns, then try again
@@ -2930,19 +2920,15 @@ def insert_chunk(
         for query, params in queries_and_params:
             try:
                 cursor = self.db.execute(query, tuple(params))
-                try: columns = [c[0] for c in cursor.description]
-                except apsw.ExecutionCompleteError: continue
-                for row in cursor:
-                    records.append(dict(zip(columns, row)))
+                cursor.row_trace = cursor_row2dict
+                records += list(cursor)
             except apsw.SQLError as e:
                 if alter and (" column" in e.args[0]):
                     # Attempt to add any missing columns, then try again
                     self.add_missing_columns(chunk)
                     cursor = self.db.execute(query, params)
-                    try: columns = [c[0] for c in cursor.description]
-                    except apsw.ExecutionCompleteError: continue
-                    for row in cursor:
-                        records.append(dict(zip(columns, row)))
+                    cursor.row_trace = cursor_row2dict
+                    records += list(cursor)
                 elif e.args[0] == "too many SQL variables":
                     first_half = chunk[: len(chunk) // 2]
                     second_half = chunk[len(chunk) // 2 :]

apswutils/utils.py

@@ -464,3 +464,8 @@ def flatten(row: dict) -> dict:
 IGNORE = object()
 SET_NULL = object()
 
+
+def cursor_row2dict(cursor, row):
+    """Converts a cursor row into a dict with columns as keys"""
+    columns = [d[0] for d in cursor.get_description()]
+    return dict(zip(columns, row))