From d1d2a8e6fa95d8daf11973f747578602d08e4962 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Wed, 26 Jan 2022 10:15:23 -0800
Subject: [PATCH] sqlite-utils bulk --batch-size option, closes #392

---
 docs/cli-reference.rst |  1 +
 docs/cli.rst           |  2 ++
 sqlite_utils/cli.py    | 16 ++++++++++++----
 tests/test_cli_bulk.py | 39 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst
index eee907e..89318d7 100644
--- a/docs/cli-reference.rst
+++ b/docs/cli-reference.rst
@@ -309,6 +309,7 @@ See :ref:`cli_bulk`.
           ' -
 
     Options:
+      --batch-size INTEGER   Commit every X records
       --flatten              Flatten nested JSON objects, so {"a": {"b": 1}} becomes
                              {"a_b": 1}
       --nl                   Expect newline-delimited JSON
diff --git a/docs/cli.rst b/docs/cli.rst
index 733eee8..3ceffce 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -1120,6 +1120,8 @@ You could insert those rows into a pre-created ``chickens`` table like so::
 
 This command takes the same options as the ``sqlite-utils insert`` command - so it defaults to expecting JSON but can accept other formats using ``--csv`` or ``--tsv`` or ``--nl`` or other options described above.
 
+By default all of the SQL queries will be executed in a single transaction. To commit every 20 records, use ``--batch-size 20``.
+
 .. _cli_insert_files:
 
 Inserting data from files
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 07a55b3..771d432 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -18,6 +18,7 @@ import csv as csv_std
 import tabulate
 from .utils import (
     _compile_code,
+    chunks,
     file_progress,
     find_spatialite,
     sqlite3,
@@ -1013,8 +1014,13 @@ def insert_upsert_implementation(
 
     # For bulk_sql= we use cursor.executemany() instead
     if bulk_sql:
-        with db.conn:
-            db.conn.cursor().executemany(bulk_sql, docs)
+        if batch_size:
+            doc_chunks = chunks(docs, batch_size)
+        else:
+            doc_chunks = [docs]
+        for doc_chunk in doc_chunks:
+            with db.conn:
+                db.conn.cursor().executemany(bulk_sql, doc_chunk)
         return
 
     try:
@@ -1256,12 +1262,14 @@ def upsert(
 )
 @click.argument("sql")
 @click.argument("file", type=click.File("rb"), required=True)
+@click.option("--batch-size", type=int, default=100, help="Commit every X records")
 @import_options
 @load_extension_option
 def bulk(
     path,
-    file,
     sql,
+    file,
+    batch_size,
     flatten,
     nl,
     csv,
@@ -1309,7 +1317,7 @@ def bulk(
             sniff=sniff,
             no_headers=no_headers,
             encoding=encoding,
-            batch_size=1,
+            batch_size=batch_size,
             alter=False,
             upsert=False,
             not_null=set(),
diff --git a/tests/test_cli_bulk.py b/tests/test_cli_bulk.py
index 37cf60b..c5dc9da 100644
--- a/tests/test_cli_bulk.py
+++ b/tests/test_cli_bulk.py
@@ -1,7 +1,11 @@
+from itertools import count
 from click.testing import CliRunner
 from sqlite_utils import cli, Database
 import pathlib
 import pytest
+import subprocess
+import sys
+import time
 
 
 @pytest.fixture
@@ -40,6 +44,41 @@ def test_cli_bulk(test_db_and_path):
     ] == list(db["example"].rows)
 
 
+def test_cli_bulk_batch_size(test_db_and_path):
+    db, db_path = test_db_and_path
+    proc = subprocess.Popen(
+        [
+            sys.executable,
+            "-m",
+            "sqlite_utils",
+            "bulk",
+            db_path,
+            "insert into example (id, name) values (:id, :name)",
+            "-",
+            "--nl",
+            "--batch-size",
+            "2",
+        ],
+        stdin=subprocess.PIPE,
+        stdout=sys.stdout,
+    )
+    # Writing one record should not commit
+    proc.stdin.write(b'{"id": 3, "name": "Three"}\n\n')
+    proc.stdin.flush()
+    time.sleep(1)
+    assert db["example"].count == 2
+
+    # Writing another should trigger a commit:
+    proc.stdin.write(b'{"id": 4, "name": "Four"}\n\n')
+    proc.stdin.flush()
+    time.sleep(1)
+    assert db["example"].count == 4
+
+    proc.stdin.close()
+    proc.wait()
+    assert proc.returncode == 0
+
+
 def test_cli_bulk_error(test_db_and_path):
     _, db_path = test_db_and_path
     result = CliRunner().invoke(