diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst index c53d64222..67e2878c4 100644 --- a/docs/cli-reference.rst +++ b/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. diff --git a/docs/cli.rst b/docs/cli.rst index 78c33b8ae..d159d8aba 100644 --- a/docs/cli.rst +++ b/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 ` 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 ` 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 diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index c23090283..46e476c82 100644 --- a/sqlite_utils/cli.py +++ b/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)) diff --git a/tests/test_cli_extract.py b/tests/test_cli_extract.py new file mode 100644 index 000000000..e646df10e --- /dev/null +++ b/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 == []