diff --git a/docs/cli.rst b/docs/cli.rst index 48cb7c1..ee25106 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -21,6 +21,43 @@ If you just want to see the FTS4 tables, you can use ``--fts4`` (or ``--fts5`` f $ sqlite-utils table_names --fts4 docs.db docs_fts +Inserting data +============== + +If you have data as JSON, you can use ``sqlite-utils insert tablename`` to insert it into a database. The table will be created with the correct (automatically detected) columns if it does not already exist. + +You can pass in a single JSON object or a list of JSON objects, either as a filename or piped directly to standard-in. + +Here's the simplest possible example:: + + $ echo '{"name": "Cleo", "age": 4}' | sqlite-utils insert dogs.db dogs - + +To specify a column as the primary key, use ``--pk=column_name``. + +If you feed it a JSON list it will insert multiple records. For example, if ``dogs.json`` looks like this:: + + [ + { + "id": 1, + "name": "Cleo", + "age": 4 + }, + { + "id": 2, + "name": "Pancakes", + "age": 2 + }, + { + "id": 3, + "name": "Toby", + "age": 6 + } + ] + +You can import all three records into an automatically created ``dogs`` table and set the ``id`` column as the primary key like so:: + + $ sqlite-utils insert dogs.db dogs.json --pk=id + Vacuum ====== diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 0025dd4..523aad8 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -1,5 +1,6 @@ import click import sqlite_utils +import json @click.group() @@ -55,3 +56,20 @@ def optimize(path, no_vacuum): db[table].optimize() if not no_vacuum: db.vacuum() + + +@cli.command() +@click.argument( + "path", + type=click.Path(file_okay=True, dir_okay=False, allow_dash=False), + required=True, +) +@click.argument("table") +@click.argument("json_file", type=click.File(), required=True) +@click.option("--pk", help="Column to use as the primary key, e.g. id") +def insert(path, table, json_file, pk): + db = sqlite_utils.Database(path) + docs = json.load(json_file) + if isinstance(docs, dict): + docs = [docs] + db[table].insert_all(docs, pk=pk) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index d9ff71e..773ac7a 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -128,6 +128,10 @@ class Table: ).fetchall() return [Column(*row) for row in rows] + @property + def pks(self): + return [column.name for column in self.columns if column.is_pk] + @property def foreign_keys(self): fks = [] diff --git a/tests/test_cli.py b/tests/test_cli.py index 429a6fe..541a7fe 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,5 +1,6 @@ from sqlite_utils import cli, Database from click.testing import CliRunner +import json import os import pytest import sqlite3 @@ -65,3 +66,47 @@ def test_optimize(db_path): # Sanity check that --no-vacuum doesn't throw errors: result = CliRunner().invoke(cli.cli, ["optimize", "--no-vacuum", db_path]) assert 0 == result.exit_code + + +def test_insert_simple(tmpdir): + json_path = str(tmpdir / "dog.json") + db_path = str(tmpdir / "dogs.db") + open(json_path, "w").write(json.dumps({"name": "Cleo", "age": 4})) + result = CliRunner().invoke(cli.cli, ["insert", db_path, "dogs", json_path]) + assert 0 == result.exit_code + assert [{"age": 4, "name": "Cleo"}] == Database(db_path).execute_returning_dicts( + "select * from dogs" + ) + db = Database(db_path) + assert ["dogs"] == db.table_names() + assert [] == db["dogs"].indexes + + +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})) + result = CliRunner().invoke( + cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id"] + ) + assert 0 == result.exit_code + assert [{"id": 1, "age": 4, "name": "Cleo"}] == Database( + db_path + ).execute_returning_dicts("select * from dogs") + db = Database(db_path) + assert ["id"] == db["dogs"].pks + + +def test_insert_multiple_with_primary_key(db_path, tmpdir): + json_path = str(tmpdir / "dogs.json") + dogs = [{"id": i, "name": "Cleo {}".format(i), "age": i + 3} for i in range(1, 21)] + open(json_path, "w").write(json.dumps(dogs)) + open("/tmp/dogs.json", "w").write(json.dumps(dogs, indent=4)) + result = CliRunner().invoke( + cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id"] + ) + assert 0 == result.exit_code + assert dogs == Database(db_path).execute_returning_dicts( + "select * from dogs order by id" + ) + db = Database(db_path) + assert ["id"] == db["dogs"].pks