From f67327abf0a9f018e1764660e190c5bbf9556ec2 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Mon, 9 Aug 2021 14:44:03 -0700 Subject: [PATCH] sqlite-utils insert --flatten option, closes #310 --- docs/cli.rst | 43 +++++++++++++++++++++++++++++++++++ sqlite_utils/cli.py | 19 ++++++++++++++++ tests/test_cli.py | 55 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 117 insertions(+) diff --git a/docs/cli.rst b/docs/cli.rst index e766d73..28bd3f7 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -722,6 +722,49 @@ This also means you pipe ``sqlite-utils`` together to easily create a new SQLite 207368,920 Kirkham St,37.760210314285,-122.47073935813 188702,1501 Evans Ave,37.7422086702947,-122.387293152263 +.. _cli_inserting_data_flatten: + +Flattening nested JSON objects +------------------------------ + +``sqlite-utils insert`` expects incoming data to consist of an array of JSON objects, where the top-level keys of each object will become columns in the created database table. + +If your data is nested you can use the `--flatten` object to create columns that are derived from the nested data. + +Consider this example document, in a file called ``log.json``:: + + { + "httpRequest": { + "latency": "0.112114537s", + "requestMethod": "GET", + "requestSize": "534", + "status": 200 + }, + "insertId": "6111722f000b5b4c4d4071e2", + "labels": { + "service": "datasette-io" + } + } + +Inserting this into a table using ``sqlite-utils insert logs.db log log.json`` will create a table with the following schema:: + + CREATE TABLE [logs] ( + [httpRequest] TEXT, + [insertId] TEXT, + [labels] TEXT + ); + +With the ``--flatten`` option columns will be created using ``topkey_nextkey`` column names - so running ``sqlite-utils insert logs.db log log.json --flatten`` will create the following schema instead:: + + CREATE TABLE [logs] ( + [httpRequest_latency] TEXT, + [httpRequest_requestMethod] TEXT, + [httpRequest_requestSize] TEXT, + [httpRequest_status] INTEGER, + [insertId] TEXT, + [labels_service] TEXT + ); + .. _cli_insert_csv_tsv: Inserting CSV or TSV data diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index c4f8501..2f32d0a 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -643,6 +643,7 @@ def insert_upsert_options(fn): "--pk", help="Columns to use as the primary key, e.g. id", multiple=True ), click.option("--nl", is_flag=True, help="Expect newline-delimited JSON"), + click.option("--flatten", is_flag=True, help="Flatten nested JSON objets"), click.option("-c", "--csv", is_flag=True, help="Expect CSV"), click.option("--tsv", is_flag=True, help="Expect TSV"), click.option("--delimiter", help="Delimiter to use for CSV files"), @@ -697,6 +698,7 @@ def insert_upsert_implementation( json_file, pk, nl, + flatten, csv, tsv, delimiter, @@ -722,6 +724,8 @@ def insert_upsert_implementation( csv = True if (nl + csv + tsv) >= 2: raise click.ClickException("Use just one of --nl, --csv or --tsv") + if (csv or tsv) and flatten: + raise click.ClickException("--flatten cannot be used with --csv or --tsv") if encoding and not (csv or tsv): raise click.ClickException("--encoding must be used with --csv or --tsv") if pk and len(pk) == 1: @@ -766,6 +770,8 @@ def insert_upsert_implementation( raise click.ClickException( "Invalid JSON - use --csv for CSV or --tsv for TSV files" ) + if flatten: + docs = (dict(_flatten(doc)) for doc in docs) extra_kwargs = {"ignore": ignore, "replace": replace, "truncate": truncate} if not_null: @@ -790,6 +796,15 @@ def insert_upsert_implementation( db[table].transform(types=tracker.types) +def _flatten(d): + for key, value in d.items(): + if isinstance(value, dict): + for key2, value2 in _flatten(value): + yield key + "_" + key2, value2 + else: + yield key, value + + @cli.command() @insert_upsert_options @click.option( @@ -813,6 +828,7 @@ def insert( json_file, pk, nl, + flatten, csv, tsv, delimiter, @@ -844,6 +860,7 @@ def insert( json_file, pk, nl, + flatten, csv, tsv, delimiter, @@ -875,6 +892,7 @@ def upsert( json_file, pk, nl, + flatten, csv, tsv, batch_size, @@ -902,6 +920,7 @@ def upsert( json_file, pk, nl, + flatten, csv, tsv, delimiter, diff --git a/tests/test_cli.py b/tests/test_cli.py index e84453a..e6f6683 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -648,6 +648,34 @@ def test_insert_invalid_json_error(tmpdir): ) +def test_insert_json_flatten(tmpdir): + db_path = str(tmpdir / "flat.db") + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "items", "-", "--flatten"], + input=json.dumps({"nested": {"data": 4}}), + ) + assert result.exit_code == 0 + assert list(Database(db_path).query("select * from items")) == [{"nested_data": 4}] + + +def test_insert_json_flatten_nl(tmpdir): + db_path = str(tmpdir / "flat.db") + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "items", "-", "--flatten", "--nl"], + input="\n".join( + json.dumps(item) + for item in [{"nested": {"data": 4}}, {"nested": {"other": 3}}] + ), + ) + assert result.exit_code == 0 + assert list(Database(db_path).query("select * from items")) == [ + {"nested_data": 4, "nested_other": None}, + {"nested_data": None, "nested_other": 3}, + ] + + def test_insert_with_primary_key(db_path, tmpdir): json_path = str(tmpdir / "dog.json") open(json_path, "w").write(json.dumps({"id": 1, "name": "Cleo", "age": 4})) @@ -1197,6 +1225,21 @@ def test_upsert(db_path, tmpdir): ] +def test_upsert_flatten(tmpdir): + db_path = str(tmpdir / "flat.db") + db = Database(db_path) + db["upsert_me"].insert({"id": 1, "name": "Example"}, pk="id") + result = CliRunner().invoke( + cli.cli, + ["upsert", db_path, "upsert_me", "-", "--flatten", "--pk", "id", "--alter"], + input=json.dumps({"id": 1, "nested": {"two": 2}}), + ) + assert result.exit_code == 0 + assert list(db.query("select * from upsert_me")) == [ + {"id": 1, "name": "Example", "nested_two": 2} + ] + + def test_upsert_alter(db_path, tmpdir): json_path = str(tmpdir / "dogs.json") db = Database(db_path) @@ -2249,3 +2292,15 @@ def test_insert_detect_types(tmpdir, option_or_env_var): _test() else: _test() + + +@pytest.mark.parametrize( + "input,expected", + ( + ({"foo": {"bar": 1}}, {"foo_bar": 1}), + ({"foo": {"bar": [1, 2, {"baz": 3}]}}, {"foo_bar": [1, 2, {"baz": 3}]}), + ({"foo": {"bar": 1, "baz": {"three": 3}}}, {"foo_bar": 1, "foo_baz_three": 3}), + ), +) +def test_flatten_helper(input, expected): + assert dict(cli._flatten(input)) == expected