--stop-after option, closes #561

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

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``).

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(),

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",
     (