Implemented 'sqlite-utils upsert my.db tablename --pk=x'

This commit is contained in:
Simon Willison 2019-01-24 21:20:10 -08:00
commit b30f924c49
3 changed files with 55 additions and 3 deletions

View file

@ -26,7 +26,7 @@ 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.
You can pass in a single JSON object or a list of JSON objects, either as a filename or piped directly to standard-in (by using ``-`` as the filename).
Here's the simplest possible example::
@ -56,7 +56,20 @@ If you feed it a JSON list it will insert multiple records. For example, if ``do
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
$ sqlite-utils insert dogs.db dogs dogs.json --pk=id
Upserting data
==============
Upserting works exactly like inserting, with the exception that if your data has a primary key that matches an already exsting record that record will be replaced with the new data.
After running the above ``dogs.json`` example, try running this::
$ echo '{"id": 2, "name": "Pancakes", "age": 3}' | \
sqlite-utils upsert dogs.db dogs - --pk=id
This will replace the record for id=2 (Pancakes) with a new record with an updated age.
Vacuum
======

View file

@ -68,8 +68,27 @@ def optimize(path, no_vacuum):
@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):
"Insert records from JSON file into the table, create table if it is missing"
db = sqlite_utils.Database(path)
docs = json.load(json_file)
if isinstance(docs, dict):
docs = [docs]
db[table].insert_all(docs, pk=pk)
@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 upsert(path, table, json_file, pk):
"Upsert records based on their primary key"
db = sqlite_utils.Database(path)
docs = json.load(json_file)
if isinstance(docs, dict):
docs = [docs]
db[table].upsert_all(docs, pk=pk)

View file

@ -100,7 +100,6 @@ 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"]
)
@ -110,3 +109,24 @@ def test_insert_multiple_with_primary_key(db_path, tmpdir):
)
db = Database(db_path)
assert ["id"] == db["dogs"].pks
def test_upsert(db_path, tmpdir):
test_insert_multiple_with_primary_key(db_path, tmpdir)
json_path = str(tmpdir / "upsert.json")
db = Database(db_path)
assert 20 == db["dogs"].count
upsert_dogs = [
{"id": 1, "name": "Upserted 1", "age": 4},
{"id": 2, "name": "Upserted 2", "age": 4},
{"id": 21, "name": "Fresh insert 21", "age": 6},
]
open(json_path, "w").write(json.dumps(upsert_dogs))
result = CliRunner().invoke(
cli.cli, ["upsert", db_path, "dogs", json_path, "--pk", "id"]
)
assert 0 == result.exit_code
assert 21 == db["dogs"].count
assert upsert_dogs == db.execute_returning_dicts(
"select * from dogs where id in (1, 2, 21) order by id"
)