insert-files multiple --pk support, closes #621

Author: simonw

sqlite_utils/cli.py

@@ -2594,7 +2594,7 @@ def extract(
     multiple=True,
     help="Column definitions for the table",
 )
-@click.option("--pk", type=str, help="Column to use as primary key")
+@click.option("pks", "--pk", help="Column to use as primary key", multiple=True)
 @click.option("--alter", is_flag=True, help="Alter table to add missing columns")
 @click.option("--replace", is_flag=True, help="Replace files with matching primary key")
 @click.option("--upsert", is_flag=True, help="Upsert files with matching primary key")
@@ -2611,7 +2611,7 @@ def insert_files(
     table,
     file_or_dir,
     column,
-    pk,
+    pks,
     alter,
     replace,
     upsert,
@@ -2641,8 +2641,8 @@ def insert_files(
             column = ["path:path", "content_text:content_text", "size:size"]
         else:
             column = ["path:path", "content:content", "size:size"]
-        if not pk:
-            pk = "path"
+        if not pks:
+            pks = ["path"]
 
     def yield_paths_and_relative_paths():
         for f_or_d in file_or_dir:
@@ -2712,7 +2712,11 @@ def _content_text(p):
         try:
             with db.conn:
                 db[table].insert_all(
-                    to_insert(), pk=pk, alter=alter, replace=replace, upsert=upsert
+                    to_insert(),
+                    pk=pks[0] if len(pks) == 1 else pks,
+                    alter=alter,
+                    replace=replace,
+                    upsert=upsert,
                 )
         except UnicodeDecodeErrorForPath as e:
             raise click.ClickException(

tests/test_insert_files.py

@@ -7,7 +7,14 @@
 
 
 @pytest.mark.parametrize("silent", (False, True))
-def test_insert_files(silent):
+@pytest.mark.parametrize(
+    "pk_args,expected_pks",
+    (
+        (["--pk", "path"], ["path"]),
+        (["--pk", "path", "--pk", "name"], ["path", "name"]),
+    ),
+)
+def test_insert_files(silent, pk_args, expected_pks):
     runner = CliRunner()
     with runner.isolated_filesystem():
         tmpdir = pathlib.Path(".")
@@ -42,7 +49,7 @@ def test_insert_files(silent):
             cli.cli,
             ["insert-files", db_path, "files", str(tmpdir)]
             + cols
-            + ["--pk", "path"]
+            + pk_args
             + (["--silent"] if silent else []),
             catch_exceptions=False,
         )
@@ -105,6 +112,7 @@ def test_insert_files(silent):
         for colname, expected_type in expected_types.items():
             for row in (one, two, three):
                 assert isinstance(row[colname], expected_type)
+        assert set(db["files"].pks) == set(expected_pks)
 
 
 @pytest.mark.parametrize(