<NME> chunked.py
<BEF> import itertools


class InvalidCallback(ValueError):
    pass


class _Chunker(object):
    def __init__(self, table, chunksize, callback):

    def __init__(self, table, chunksize=1000):
        self.queue = []
        self.table = table
        self.chunksize = chunksize

    def insert(self, item):
        self.queue.append(item)
        if len(self.queue) > self.chunksize:
            self.table.insert_many(self.queue)
            self.queue = []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.table.insert_many(self.queue)
        self.queue = []
        inserter(row)

    Rows will be inserted in groups of `chunksize` (defaulting to 1000). An
    optional callback can be provided that will be called before the insert.
    This callback takes one parameter which is the queue which is about to be
    inserted into the database
    """

    def __init__(self, table, chunksize=1000, callback=None):
        self.fields = set()
        super().__init__(table, chunksize, callback)

    def insert(self, item):
        self.fields.update(item.keys())
        super()._queue_add(item)

    def flush(self):
        for item in self.queue:
            for field in self.fields:
                item[field] = item.get(field)
        if self.callback is not None:
            self.callback(self.queue)
        self.table.insert_many(self.queue)
        super().flush()


class ChunkedUpdate(_Chunker):
    """Batch up update operations
    with ChunkedUpdate(my_table) as updater:
        updater(row)

    Rows will be updated in groups of `chunksize` (defaulting to 1000). An
    optional callback can be provided that will be called before the update.
    This callback takes one parameter which is the queue which is about to be
    updated into the database
    """

    def __init__(self, table, keys, chunksize=1000, callback=None):
        self.keys = keys
        super().__init__(table, chunksize, callback)

    def update(self, item):
        super()._queue_add(item)

    def flush(self):
        if self.callback is not None:
            self.callback(self.queue)
        self.queue.sort(key=dict.keys)
        for fields, items in itertools.groupby(self.queue, key=dict.keys):
            self.table.update_many(list(items), self.keys)
        super().flush()

<MSG> make sure all records have the same fields.

<DFF> @@ -10,18 +10,25 @@ class ChunkedInsert(object):
 
     def __init__(self, table, chunksize=1000):
         self.queue = []
+        self.fields = set()
         self.table = table
         self.chunksize = chunksize
 
+    def flush(self):
+        for item in self.queue:
+            for field in self.fields:
+                item[field] = item.get(field)
+        self.table.insert_many(self.queue)
+        self.queue = []
+
     def insert(self, item):
+        self.fields.update(item.keys())
         self.queue.append(item)
-        if len(self.queue) > self.chunksize:
-            self.table.insert_many(self.queue)
-            self.queue = []
+        if len(self.queue) >= self.chunksize:
+            self.flush()
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.table.insert_many(self.queue)
-        self.queue = []
+        self.flush()

{'addition_count': 12, 'commit_subject': 'make sure all records have the same fields.', 'deletion_count': 5, 'file_extension': '.py', 'lang': 'py', 'license': 'mit', 'repo_name': 'pudo/dataset'}