sqlite-utils bulk --batch-size option, closes #392

This commit is contained in:
Simon Willison 2022-01-26 10:15:23 -08:00
commit d1d2a8e6fa
4 changed files with 54 additions and 4 deletions

View file

@ -309,6 +309,7 @@ See :ref:`cli_bulk`.
' - ' -
Options: Options:
--batch-size INTEGER Commit every X records
--flatten Flatten nested JSON objects, so {"a": {"b": 1}} becomes --flatten Flatten nested JSON objects, so {"a": {"b": 1}} becomes
{"a_b": 1} {"a_b": 1}
--nl Expect newline-delimited JSON --nl Expect newline-delimited JSON

View file

@ -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. 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: .. _cli_insert_files:
Inserting data from files Inserting data from files

View file

@ -18,6 +18,7 @@ import csv as csv_std
import tabulate import tabulate
from .utils import ( from .utils import (
_compile_code, _compile_code,
chunks,
file_progress, file_progress,
find_spatialite, find_spatialite,
sqlite3, sqlite3,
@ -1013,8 +1014,13 @@ def insert_upsert_implementation(
# For bulk_sql= we use cursor.executemany() instead # For bulk_sql= we use cursor.executemany() instead
if bulk_sql: if bulk_sql:
with db.conn: if batch_size:
db.conn.cursor().executemany(bulk_sql, docs) 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 return
try: try:
@ -1256,12 +1262,14 @@ def upsert(
) )
@click.argument("sql") @click.argument("sql")
@click.argument("file", type=click.File("rb"), required=True) @click.argument("file", type=click.File("rb"), required=True)
@click.option("--batch-size", type=int, default=100, help="Commit every X records")
@import_options @import_options
@load_extension_option @load_extension_option
def bulk( def bulk(
path, path,
file,
sql, sql,
file,
batch_size,
flatten, flatten,
nl, nl,
csv, csv,
@ -1309,7 +1317,7 @@ def bulk(
sniff=sniff, sniff=sniff,
no_headers=no_headers, no_headers=no_headers,
encoding=encoding, encoding=encoding,
batch_size=1, batch_size=batch_size,
alter=False, alter=False,
upsert=False, upsert=False,
not_null=set(), not_null=set(),

View file

@ -1,7 +1,11 @@
from itertools import count
from click.testing import CliRunner from click.testing import CliRunner
from sqlite_utils import cli, Database from sqlite_utils import cli, Database
import pathlib import pathlib
import pytest import pytest
import subprocess
import sys
import time
@pytest.fixture @pytest.fixture
@ -40,6 +44,41 @@ def test_cli_bulk(test_db_and_path):
] == list(db["example"].rows) ] == 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): def test_cli_bulk_error(test_db_and_path):
_, db_path = test_db_and_path _, db_path = test_db_and_path
result = CliRunner().invoke( result = CliRunner().invoke(