Merge 7c01b8d588 into a947dc6739

docs/python-api.rst

@@ -1088,6 +1088,12 @@ The function can accept an iterator or generator of rows and will commit them ac
         "name": "Name {}".format(i),
     } for i in range(10000)), batch_size=1000)
 
+The largest batch that will actually be sent to SQLite is limited by the maximum number of SQL variables allowed in a single query, which defaults to 999. If your copy of SQLite was compiled with a higher ``SQLITE_MAX_VARIABLE_NUMBER`` you can tell ``sqlite-utils`` to use larger batches - and hence run faster - by passing ``sqlite_max_vars=`` to the ``Database()`` constructor:
+
+.. code-block:: python
+
+    db = Database("big.db", sqlite_max_vars=100_000)
+
 You can skip inserting any records that have a primary key that already exists using ``ignore=True``. This works with both ``.insert({...}, ignore=True)`` and ``.insert_all([...], ignore=True)``.
 
 You can delete all the existing rows in the table before inserting the new records using ``truncate=True``. This is useful if you want to replace the data in the table.

sqlite_utils/db.py

@@ -504,6 +504,9 @@ class Database:
     :param use_old_upsert: set to ``True`` to force the older upsert implementation. See
       :ref:`python_api_old_upsert`
     :param strict: Apply STRICT mode to all created tables (unless overridden)
+    :param sqlite_max_vars: Maximum number of SQL variables to use in a single query. Defaults
+      to ``sqlite_utils.db.SQLITE_MAX_VARS`` (999). Increase this if your SQLite was compiled
+      with a higher ``SQLITE_MAX_VARIABLE_NUMBER`` to allow larger insert batches
     """
 
     _counts_table_name = "_counts"
@@ -522,10 +525,12 @@ class Database:
         execute_plugins: bool = True,
         use_old_upsert: bool = False,
         strict: bool = False,
+        sqlite_max_vars: Optional[int] = None,
     ):
         self.memory_name = None
         self.memory = False
         self.use_old_upsert = use_old_upsert
+        self._sqlite_max_vars = sqlite_max_vars
         if not (
             (filename_or_conn is not None and (not memory and not memory_name))
             or (filename_or_conn is None and (memory or memory_name))
@@ -579,6 +584,17 @@ class Database:
             pm.hook.prepare_connection(conn=self.conn)
         self.strict = strict
 
+    @property
+    def sqlite_max_vars(self) -> int:
+        """
+        The maximum number of SQL variables to use in a single query. This is the value
+        passed as ``sqlite_max_vars=`` to the constructor, or the
+        ``sqlite_utils.db.SQLITE_MAX_VARS`` default of 999 if that was not set.
+        """
+        if self._sqlite_max_vars is not None:
+            return self._sqlite_max_vars
+        return SQLITE_MAX_VARS
+
     def __enter__(self) -> "Database":
         return self
 
@@ -4482,14 +4498,11 @@ class Table(Queryable):
             first_record = cast(Dict[str, Any], first_record)
             num_columns = len(first_record.keys())
 
-        if num_columns > SQLITE_MAX_VARS:
+        max_vars = self.db.sqlite_max_vars
-            raise ValueError(
+        if num_columns > max_vars:
-                "Rows can have a maximum of {} columns".format(SQLITE_MAX_VARS)
+            raise ValueError("Rows can have a maximum of {} columns".format(max_vars))
-            )
         batch_size = (
-            1
+            1 if num_columns == 0 else max(1, min(batch_size, max_vars // num_columns))
-            if num_columns == 0
-            else max(1, min(batch_size, SQLITE_MAX_VARS // num_columns))
         )
         self.last_rowid = None
         self.last_pk = None

tests/test_create.py

@@ -695,6 +695,39 @@ def test_bulk_insert_more_than_999_values(fresh_db):
     assert fresh_db["big"].count == 100
 
 
+def test_sqlite_max_vars_defaults_to_999():
+    # https://github.com/simonw/sqlite-utils/issues/147
+    assert Database(memory=True).sqlite_max_vars == 999
+
+
+def test_sqlite_max_vars_can_be_customized():
+    # https://github.com/simonw/sqlite-utils/issues/147
+    assert Database(memory=True, sqlite_max_vars=100000).sqlite_max_vars == 100000
+    # A raised limit should allow a bigger batch, so the same records are
+    # written using fewer INSERT statements
+    records = [{"c{}".format(i): i for i in range(5)} for _ in range(500)]
+
+    def count_inserts(sqlite_max_vars):
+        seen = []
+        db = Database(
+            memory=True,
+            sqlite_max_vars=sqlite_max_vars,
+            tracer=lambda sql, params: seen.append(sql),
+        )
+        db["t"].insert_all(records, batch_size=100000)
+        return len([sql for sql in seen if sql.strip().upper().startswith("INSERT")])
+
+    assert count_inserts(100000) == 1
+    assert count_inserts(None) > 1
+
+
+def test_error_message_uses_custom_sqlite_max_vars():
+    # https://github.com/simonw/sqlite-utils/issues/147
+    db = Database(memory=True, sqlite_max_vars=10)
+    with pytest.raises(ValueError, match="maximum of 10 columns"):
+        db["big"].insert({"c{}".format(i): i for i in range(11)})
+
+
 @pytest.mark.parametrize(
     "num_columns,should_error", ((900, False), (999, False), (1000, True))
 )