Add --extract to CLI insert and upsert

docs/cli-reference.rst

@@ -315,6 +315,7 @@ See :ref:`cli_inserting_data`, :ref:`cli_insert_csv_tsv`, :ref:`cli_insert_unstr
       --analyze                 Run ANALYZE at the end of this operation
       --load-extension TEXT     Path to SQLite extension, with optional :entrypoint
       --silent                  Do not show progress bar
+      --extract TEXT            Columns to extract to same-name lookup tables
       --strict                  Apply STRICT mode to created table
       --ignore                  Ignore records if pk already exists
       --replace                 Replace records if pk already exists
@@ -379,6 +380,7 @@ See :ref:`cli_upsert`.
       --analyze                 Run ANALYZE at the end of this operation
       --load-extension TEXT     Path to SQLite extension, with optional :entrypoint
       --silent                  Do not show progress bar
+      --extract TEXT            Columns to extract to same-name lookup tables
       --strict                  Apply STRICT mode to created table
       -h, --help                Show this message and exit.
 

docs/cli.rst

@@ -1312,6 +1312,33 @@ With the ``--flatten`` option columns will be created using ``topkey_nextkey`` c
        "labels_service" TEXT
     );
 
+.. _cli_insert_extract:
+
+Extracting values during an import
+==================================
+
+Use ``--extract COLUMN`` with ``insert`` or ``upsert`` to move repeated values into a lookup table as records are imported:
+
+.. code-block:: bash
+
+    echo '[{"id": 1, "species": "oak"}, {"id": 2, "species": "oak"}]' \
+      | sqlite-utils insert trees.db trees - --pk id --extract species
+
+This creates a ``species`` table with ``id`` and ``value`` columns and a unique index on ``value``. Both input records reference the same lookup row. In a newly created ``trees`` table, the ``species`` column is an integer foreign key referencing ``species.id``. Extracted ``null`` values stay ``null`` and do not create lookup rows.
+
+Repeat the option to extract multiple columns into separate lookup tables:
+
+.. code-block:: bash
+
+    sqlite-utils insert trees.db trees trees.csv --csv \
+      --extract species --extract country
+
+The option also works with TSV, newline-delimited JSON, converted rows and ``--code``. CSV/TSV type detection still applies to other columns; extracted columns in a newly created table keep their integer foreign key type. Subsequent imports and upserts reuse existing lookup rows.
+
+This uses the Python API's :ref:`extracts= mechanism <python_api_extracts>` and keeps its behavior for existing tables: it does not change an existing column's type or add a foreign key constraint to that column. Start with a new table to create the foreign key schema, or use :ref:`sqlite-utils extract <cli_extract>` to normalize values already stored in a table.
+
+Without ``--extract``, imported values are stored directly as before.
+
 .. _cli_insert_csv_tsv:
 
 Inserting CSV or TSV data

sqlite_utils/cli.py

@@ -1044,6 +1044,12 @@ def inner(fn):
                 ),
                 load_extension_option,
                 click.option("--silent", is_flag=True, help="Do not show progress bar"),
+                click.option(
+                    "--extract",
+                    "extracts",
+                    multiple=True,
+                    help="Columns to extract to same-name lookup tables",
+                ),
                 click.option(
                     "--strict",
                     is_flag=True,
@@ -1095,6 +1101,7 @@ def insert_upsert_implementation(
     functions=None,
     strict=False,
     code=None,
+    extracts=None,
 ):
     db = sqlite_utils.Database(path)
     _register_db_for_cleanup(db)
@@ -1118,6 +1125,8 @@ def _insert_docs(docs, tracker=None):
             extra_kwargs["columns"] = column_type_overrides
         if upsert:
             extra_kwargs["upsert"] = upsert
+        if extracts:
+            extra_kwargs["extracts"] = extracts
 
         # docs should all be dictionaries
         docs = (verify_is_dict(doc) for doc in docs)
@@ -1176,6 +1185,9 @@ def _insert_docs(docs, tracker=None):
         ):
             detected_types = tracker.types
             detected_types.update(column_type_overrides)
+            # Extracted values have been replaced by integer foreign keys.
+            for column in extracts or ():
+                detected_types.pop(column, None)
             db.table(table).transform(types=detected_types)
 
     if code is not None:
@@ -1386,6 +1398,7 @@ def insert(
     default,
     types,
     strict,
+    extracts,
 ):
     """
     Insert records from FILE into a table, creating the table if it
@@ -1481,6 +1494,7 @@ def rows():
             types=types,
             strict=strict,
             code=code,
+            extracts=extracts,
         )
     except UnicodeDecodeError as ex:
         raise click.ClickException(UNICODE_ERROR.format(ex))
@@ -1519,6 +1533,7 @@ def upsert(
     load_extension,
     silent,
     strict,
+    extracts,
 ):
     """
     Upsert records based on their primary key. Works like 'insert' but if
@@ -1571,6 +1586,7 @@ def upsert(
             silent=silent,
             strict=strict,
             code=code,
+            extracts=extracts,
         )
     except UnicodeDecodeError as ex:
         raise click.ClickException(UNICODE_ERROR.format(ex))

tests/test_cli_extract.py

@@ -0,0 +1,243 @@
+import csv
+import io
+import json
+
+import pytest
+from click.testing import CliRunner
+
+from sqlite_utils import Database, cli
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+@pytest.mark.parametrize("format", ["json", "nl", "csv", "tsv"])
+@pytest.mark.parametrize("batch_size", [1, 100])
+def test_cli_extract(tmp_path, command, format, batch_size):
+    path = str(tmp_path / "trees.db")
+    rows = [
+        {"id": 1, "species": "oak", "region": "west", "height": 1.5},
+        {"id": 2, "species": "ash", "region": "west", "height": 2.5},
+        {"id": 3, "species": "oak", "region": "east", "height": 3.5},
+    ]
+    if format == "json":
+        input = json.dumps(rows)
+    elif format == "nl":
+        input = "\n".join(json.dumps(row) for row in rows)
+    else:
+        output = io.StringIO()
+        writer = csv.DictWriter(
+            output,
+            fieldnames=list(rows[0]),
+            delimiter="\t" if format == "tsv" else ",",
+        )
+        writer.writeheader()
+        writer.writerows(rows)
+        input = output.getvalue()
+    args = [
+        command,
+        path,
+        "trees",
+        "-",
+        "--pk",
+        "id",
+        "--extract",
+        "species",
+        "--extract",
+        "region",
+        "--batch-size",
+        str(batch_size),
+    ]
+    if format != "json":
+        args.append("--" + format)
+    result = CliRunner().invoke(cli.cli, args, input=input)
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["trees"].rows) == [
+        {"id": 1, "species": 1, "region": 1, "height": 1.5},
+        {"id": 2, "species": 2, "region": 1, "height": 2.5},
+        {"id": 3, "species": 1, "region": 2, "height": 3.5},
+    ]
+    assert list(db["species"].rows) == [
+        {"id": 1, "value": "oak"},
+        {"id": 2, "value": "ash"},
+    ]
+    assert list(db["region"].rows) == [
+        {"id": 1, "value": "west"},
+        {"id": 2, "value": "east"},
+    ]
+    assert db["trees"].columns_dict == {
+        "id": int,
+        "species": int,
+        "region": int,
+        "height": float,
+    }
+    assert {
+        (fk.column, fk.other_table, fk.other_column) for fk in db["trees"].foreign_keys
+    } == {
+        ("species", "species", "id"),
+        ("region", "region", "id"),
+    }
+    assert list(db.query("pragma foreign_key_check")) == []
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_extract_code(tmp_path, command):
+    path = str(tmp_path / "code.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [
+            command,
+            path,
+            "trees",
+            "--code",
+            'rows = [{"id": 1, "species": "oak"}]',
+            "--pk",
+            "id",
+            "--extract",
+            "species",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["trees"].rows) == [{"id": 1, "species": 1}]
+    assert list(db["species"].rows) == [{"id": 1, "value": "oak"}]
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+@pytest.mark.parametrize(
+    "extra_args",
+    [[], ["--strict"], ["--no-detect-types"], ["--type", "species", "TEXT"]],
+)
+def test_cli_extract_csv_types(tmp_path, command, extra_args):
+    path = str(tmp_path / "types.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [command, path, "trees", "-", "--pk", "id", "--csv", "--extract", "species"]
+        + extra_args,
+        input="id,species,height\n1,oak,1.5\n",
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert db["trees"].columns_dict["species"] is int
+    assert list(db["trees"].rows)[0]["species"] == 1
+    assert list(db.query("pragma foreign_key_check")) == []
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_extract_nulls(tmp_path, command):
+    path = str(tmp_path / "nulls.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [command, path, "trees", "-", "--pk", "id", "--extract", "species"],
+        input=json.dumps(
+            [
+                {"id": 1, "species": None},
+                {"id": 2, "species": "oak"},
+                {"id": 3, "species": None},
+            ]
+        ),
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["species"].rows) == [{"id": 1, "value": "oak"}]
+    assert list(db["trees"].rows) == [
+        {"id": 1, "species": None},
+        {"id": 2, "species": 1},
+        {"id": 3, "species": None},
+    ]
+
+
+def test_cli_extract_upsert(tmp_path):
+    path = str(tmp_path / "upsert.db")
+    for command, rows in [
+        ("insert", [{"id": 1, "species": "oak"}, {"id": 2, "species": "ash"}]),
+        ("upsert", [{"id": 1, "species": "ash"}, {"id": 3, "species": "oak"}]),
+    ]:
+        result = CliRunner().invoke(
+            cli.cli,
+            [command, path, "trees", "-", "--pk", "id", "--extract", "species"],
+            input=json.dumps(rows),
+        )
+        assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["species"].rows) == [
+        {"id": 1, "value": "oak"},
+        {"id": 2, "value": "ash"},
+    ]
+    assert list(db["trees"].rows) == [
+        {"id": 1, "species": 2},
+        {"id": 2, "species": 2},
+        {"id": 3, "species": 1},
+    ]
+
+
+@pytest.mark.parametrize("column", ['tree "species"', "品種", ""])
+def test_cli_extract_identifiers(tmp_path, column):
+    path = str(tmp_path / "identifiers.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        ["insert", path, "trees", "-", "--extract", column],
+        input=json.dumps({column: "oak"}),
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert list(db["trees"].rows) == [{column: 1}]
+    assert list(db[column].rows) == [{"id": 1, "value": "oak"}]
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_extract_errors(tmp_path, command):
+    path = str(tmp_path / "invalid.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [command, path, "trees", "-", "--pk", "id", "--extract", "species"],
+        input="invalid JSON",
+    )
+    assert result.exit_code == 1
+    assert "Invalid JSON" in result.output
+    assert Database(path).table_names() == []
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_extract_requires_column(tmp_path, command):
+    path = tmp_path / "missing.db"
+    result = CliRunner().invoke(
+        cli.cli, [command, str(path), "trees", "-", "--extract"], input="{}"
+    )
+    assert result.exit_code == 2
+    assert "requires an argument" in result.output
+    assert not path.exists()
+
+
+def test_cli_extract_existing_schema(tmp_path):
+    path = str(tmp_path / "existing.db")
+    db = Database(path)
+    db["trees"].create({"id": int, "species": int}, pk="id")
+    schema = db["trees"].schema
+    result = CliRunner().invoke(
+        cli.cli,
+        ["insert", path, "trees", "-", "--extract", "species"],
+        input='{"id": 1, "species": "oak"}',
+    )
+    assert result.exit_code == 0, result.output
+    assert db["trees"].schema == schema
+    assert db["trees"].foreign_keys == []
+    assert list(db["trees"].rows) == [{"id": 1, "species": 1}]
+    assert list(db["species"].rows) == [{"id": 1, "value": "oak"}]
+
+
+@pytest.mark.parametrize("command", ["insert", "upsert"])
+def test_cli_no_extract_default(tmp_path, command):
+    path = str(tmp_path / "default.db")
+    result = CliRunner().invoke(
+        cli.cli,
+        [command, path, "trees", "-", "--pk", "id"],
+        input='[{"id": 1, "species": "oak"}, {"id": 2, "species": "oak"}]',
+    )
+    assert result.exit_code == 0, result.output
+    db = Database(path)
+    assert db.table_names() == ["trees"]
+    assert list(db["trees"].rows) == [
+        {"id": 1, "species": "oak"},
+        {"id": 2, "species": "oak"},
+    ]
+    assert db["trees"].foreign_keys == []