From 8c739558f710961464d6fb3f4ef8b850dc172ab5 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 27 Jun 2023 11:50:04 -0700
Subject: [PATCH] --stop-after option, closes #561

---
 docs/cli-reference.rst   |  2 ++
 docs/cli.rst             |  6 ++++++
 sqlite_utils/cli.py      | 10 ++++++++++
 tests/test_cli_insert.py | 25 +++++++++++++++++++++++++
 4 files changed, 43 insertions(+)

diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst
index e38d08d..980a66b 100644
--- a/docs/cli-reference.rst
+++ b/docs/cli-reference.rst
@@ -276,6 +276,7 @@ See :ref:`cli_inserting_data`, :ref:`cli_insert_csv_tsv`, :ref:`cli_insert_unstr
       --no-headers              CSV file has no header row
       --encoding TEXT           Character encoding for input, defaults to utf-8
       --batch-size INTEGER      Commit every X records
+      --stop-after INTEGER      Stop after X records
       --alter                   Alter existing table to add any missing columns
       --not-null TEXT           Columns that should be created as NOT NULL
       --default <TEXT TEXT>...  Default value that should be set for a column
@@ -330,6 +331,7 @@ See :ref:`cli_upsert`.
       --no-headers              CSV file has no header row
       --encoding TEXT           Character encoding for input, defaults to utf-8
       --batch-size INTEGER      Commit every X records
+      --stop-after INTEGER      Stop after X records
       --alter                   Alter existing table to add any missing columns
       --not-null TEXT           Columns that should be created as NOT NULL
       --default <TEXT TEXT>...  Default value that should be set for a column
diff --git a/docs/cli.rst b/docs/cli.rst
index a121674..77ef0ee 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -1237,6 +1237,12 @@ Data is expected to be encoded as Unicode UTF-8. If your data is an another char
 
     sqlite-utils insert dogs.db dogs dogs.tsv --tsv --encoding=latin-1
 
+To stop inserting after a specified number of records - useful for getting a faster preview of a large file - use the ``--stop-after`` option:
+
+.. code-block:: bash
+
+    sqlite-utils insert dogs.db dogs dogs.csv --csv --stop-after=10
+
 A progress bar is displayed when inserting data from a file. You can hide the progress bar using the ``--silent`` option.
 
 By default every column inserted from a CSV or TSV file will be of type ``TEXT``. To automatically detect column types - resulting in a mix of ``TEXT``, ``INTEGER`` and ``FLOAT`` columns, use the ``--detect-types`` option (or its shortcut ``-d``).
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index bbcc888..6ebe407 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -868,6 +868,7 @@ def insert_upsert_options(*, require_pk=False):
                 click.option(
                     "--batch-size", type=int, default=100, help="Commit every X records"
                 ),
+                click.option("--stop-after", type=int, help="Stop after X records"),
                 click.option(
                     "--alter",
                     is_flag=True,
@@ -925,6 +926,7 @@ def insert_upsert_implementation(
     no_headers,
     encoding,
     batch_size,
+    stop_after,
     alter,
     upsert,
     ignore=False,
@@ -1012,6 +1014,9 @@ def insert_upsert_implementation(
             if flatten:
                 docs = (_flatten(doc) for doc in docs)
 
+    if stop_after:
+        docs = itertools.islice(docs, stop_after)
+
     if convert:
         variable = "row"
         if lines:
@@ -1146,6 +1151,7 @@ def insert(
     no_headers,
     encoding,
     batch_size,
+    stop_after,
     alter,
     detect_types,
     analyze,
@@ -1221,6 +1227,7 @@ def insert(
             no_headers,
             encoding,
             batch_size,
+            stop_after,
             alter=alter,
             upsert=False,
             ignore=ignore,
@@ -1253,6 +1260,7 @@ def upsert(
     convert,
     imports,
     batch_size,
+    stop_after,
     delimiter,
     quotechar,
     sniff,
@@ -1299,6 +1307,7 @@ def upsert(
             no_headers,
             encoding,
             batch_size,
+            stop_after,
             alter=alter,
             upsert=True,
             not_null=not_null,
@@ -1380,6 +1389,7 @@ def bulk(
             no_headers=no_headers,
             encoding=encoding,
             batch_size=batch_size,
+            stop_after=None,
             alter=False,
             upsert=False,
             not_null=set(),
diff --git a/tests/test_cli_insert.py b/tests/test_cli_insert.py
index f403328..ca261aa 100644
--- a/tests/test_cli_insert.py
+++ b/tests/test_cli_insert.py
@@ -229,6 +229,31 @@ def test_insert_csv_tsv(content, options, db_path, tmpdir):
     assert [{"foo": "1", "bar": "2", "baz": "cat,dog"}] == list(db["data"].rows)
 
 
+@pytest.mark.parametrize(
+    "input,args",
+    (
+        (
+            json.dumps(
+                [{"name": "One"}, {"name": "Two"}, {"name": "Three"}, {"name": "Four"}]
+            ),
+            [],
+        ),
+        ("name\nOne\nTwo\nThree\nFour\n", ["--csv"]),
+    ),
+)
+def test_insert_stop_after(tmpdir, input, args):
+    db_path = str(tmpdir / "data.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        ["insert", db_path, "rows", "-", "--stop-after", "2"] + args,
+        input=input,
+    )
+    assert 0 == result.exit_code
+    assert [{"name": "One"}, {"name": "Two"}] == list(
+        Database(db_path).query("select * from rows")
+    )
+
+
 @pytest.mark.parametrize(
     "options",
     (