insert_all() now accepts generator, closes #7

docs/python-api.rst

@@ -80,8 +80,6 @@ This will automatically create a new table called "dogs" with the following sche
         is_good_dog INTEGER
     )
 
-The column types are automatically derived from the types of the incoming data.
-
 You can also specify a primary key by passing the ``pk=`` parameter to the ``.insert()`` call. This will only be obeyed if the record being inserted causes the table to be created:
 
 .. code-block:: python
@@ -113,7 +111,7 @@ You don't need to pass all of the columns to the ``column_order`` parameter. If
 Bulk inserts
 ============
 
-If you have more than one record to insert, the ``insert_all()`` method is a much more efficient way of inserting them. Just like ``insert()`` it will automatically detect the columns that should be created, but it will inspect the first 100 items to help decide what those column types should be.
+If you have more than one record to insert, the ``insert_all()`` method is a much more efficient way of inserting them. Just like ``insert()`` it will automatically detect the columns that should be created, but it will inspect the first batch of 100 items to help decide what those column types should be.
 
 Use it like this:
 
@@ -133,6 +131,17 @@ Use it like this:
         "is_good_dog": True,
     }], pk="id", column_order=("id", "twitter", "name"))
 
+The column types used in the ``CREATE TABLE`` statement are automatically derived from the types of data in that first batch of rows. Any additional or missing columns in subsequent batches will be ignored.
+
+The function can accept an iterator or generator of rows and will commit them according to the batch size. The default batch size is 100, but you can specify a different size using the ``batch_size`` parameter:
+
+.. code-block:: python
+
+    db["big_table"].insert_all(({
+        "id": 1,
+        "name": "Name {}".format(i),
+    } for i in range(10000)), batch_size=1000)
+
 Upserting data
 ==============
 

sqlite_utils/db.py

@@ -1,6 +1,7 @@
 import sqlite3
 from collections import namedtuple
 import datetime
+import itertools
 import json
 import pathlib
 
@@ -345,18 +346,24 @@ class Table:
         that it creates (if table does not exist) has columns for ALL of that
         data
         """
-        if not self.exists:
-            self.create(
-                self.detect_column_types(records),
-                pk,
-                foreign_keys,
-                column_order=column_order,
-            )
-        all_columns = set()
-        for record in records:
-            all_columns.update(record.keys())
-        all_columns = list(sorted(all_columns))
+        all_columns = None
+        first = True
         for chunk in chunks(records, batch_size):
+            chunk = list(chunk)
+            if first:
+                if not self.exists:
+                    # Use the first batch to derive the table names
+                    self.create(
+                        self.detect_column_types(chunk),
+                        pk,
+                        foreign_keys,
+                        column_order=column_order,
+                    )
+                all_columns = set()
+                for record in chunk:
+                    all_columns.update(record.keys())
+                all_columns = list(sorted(all_columns))
+            first = False
             sql = """
                 INSERT {upsert} INTO [{table}] ({columns}) VALUES {rows};
             """.format(
@@ -402,8 +409,9 @@ class Table:
 
 
 def chunks(sequence, size):
-    for i in range(0, len(sequence), size):
-        yield sequence[i : i + size]
+    iterator = iter(sequence)
+    for item in iterator:
+        yield itertools.chain([item], itertools.islice(iterator, size - 1))
 
 
 def jsonify_if_needed(value):

tests/test_create.py

@@ -172,6 +172,25 @@ def test_insert_dictionaries_and_lists_as_json(fresh_db, data_structure):
     assert data_structure == json.loads(row[1])
 
 
+def test_insert_thousands_using_generator(fresh_db):
+    fresh_db["test"].insert_all(
+        {"i": i, "word": "word_{}".format(i)} for i in range(10000)
+    )
+    assert [{"name": "i", "type": "INTEGER"}, {"name": "word", "type": "TEXT"}] == [
+        {"name": col.name, "type": col.type} for col in fresh_db["test"].columns
+    ]
+    assert 10000 == fresh_db["test"].count
+
+
+def test_insert_thousands_ignores_extra_columns_after_first_100(fresh_db):
+    fresh_db["test"].insert_all(
+        [{"i": i, "word": "word_{}".format(i)} for i in range(100)]
+        + [{"i": 101, "extra": "This extra column should cause an exception"}]
+    )
+    rows = fresh_db.execute_returning_dicts("select * from test where i = 101")
+    assert [{"i": 101, "word": None}] == rows
+
+
 def test_create_view(fresh_db):
     fresh_db["data"].insert({"foo": "foo", "bar": "bar"})
     fresh_db.create_view("bar", "select bar from data")