diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 53f8ee0..b7de419 100644 --- a/sqlite_utils/cli.py +++ b/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 insert_files( 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( diff --git a/tests/test_insert_files.py b/tests/test_insert_files.py index d7af28e..88e49a8 100644 --- a/tests/test_insert_files.py +++ b/tests/test_insert_files.py @@ -7,7 +7,14 @@ import sys @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(