CLI now accepts multiple --pk for compound primary keys

This commit is contained in:
Simon Willison 2019-07-14 20:57:57 -07:00
commit f47b5a9e19
2 changed files with 25 additions and 1 deletions

View file

@ -302,7 +302,7 @@ def insert_upsert_options(fn):
),
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"),
click.option("--pk", help="Column to use as the primary key, e.g. id", multiple=True),
click.option("--nl", is_flag=True, help="Expect newline-delimited JSON"),
click.option("-c", "--csv", is_flag=True, help="Expect CSV"),
click.option(
@ -348,6 +348,8 @@ def insert_upsert_implementation(
if nl and csv:
click.echo("Use just one of --nl and --csv", err=True)
return
if pk and len(pk) == 1:
pk = pk[0]
if csv:
reader = csv_std.reader(json_file)
headers = next(reader)

View file

@ -427,6 +427,28 @@ def test_insert_multiple_with_primary_key(db_path, tmpdir):
assert ["id"] == db["dogs"].pks
def test_insert_multiple_with_compound_primary_key(db_path, tmpdir):
json_path = str(tmpdir / "dogs.json")
dogs = [{"breed": "mixed", "id": i, "name": "Cleo {}".format(i), "age": i + 3} for i in range(1, 21)]
open(json_path, "w").write(json.dumps(dogs))
result = CliRunner().invoke(
cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id", "--pk", "breed"]
)
assert 0 == result.exit_code
db = Database(db_path)
assert dogs == db.execute_returning_dicts("select * from dogs order by breed, id")
assert {"breed", "id"} == set(db["dogs"].pks)
assert (
"CREATE TABLE [dogs] (\n"
" [breed] TEXT,\n"
" [id] INTEGER,\n"
" [name] TEXT,\n"
" [age] INTEGER,\n"
" PRIMARY KEY ([id], [breed])\n"
")"
) == db["dogs"].schema
def test_insert_not_null_default(db_path, tmpdir):
json_path = str(tmpdir / "dogs.json")
dogs = [