sqlite-utils insert --flatten option, closes #310

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

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,

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