From f47b5a9e1954bc6aefcca4a822b24f289c25f1af Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 14 Jul 2019 20:57:57 -0700
Subject: [PATCH] CLI now accepts multiple --pk for compound primary keys

---
 sqlite_utils/cli.py |  4 +++-
 tests/test_cli.py   | 22 ++++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 1c8661f..34228e4 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -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)
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 47afbb0..5526d4d 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -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 = [