remote/sqlite-utils
1
0
Fork 0
mirror of https://github.com/simonw/sqlite-utils.git synced 2026-09-10 10:34:10 +02:00
Code Issues Projects Releases Packages Wiki Activity

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

Browse source
This commit is contained in:
Simon Willison 2024-03-16 17:33:31 -07:00
commit 1d050dcdc7
2 changed files with 19 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

14
sqlite_utils/cli.py
View file

@ -2594,7 +2594,7 @@ def extract(
multiple=True, multiple=True,
help="Column definitions for the table", 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("--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("--replace", is_flag=True, help="Replace files with matching primary key")
@click.option("--upsert", is_flag=True, help="Upsert 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, table,
file_or_dir, file_or_dir,
column, column,
pk, pks,
alter, alter,
replace, replace,
upsert, upsert,
@ -2641,8 +2641,8 @@ def insert_files(
column = ["path:path", "content_text:content_text", "size:size"] column = ["path:path", "content_text:content_text", "size:size"]
else: else:
column = ["path:path", "content:content", "size:size"] column = ["path:path", "content:content", "size:size"]
if not pk: if not pks:
pk = "path" pks = ["path"]
def yield_paths_and_relative_paths(): def yield_paths_and_relative_paths():
for f_or_d in file_or_dir: for f_or_d in file_or_dir:
@ -2712,7 +2712,11 @@ def insert_files(
try: try:
with db.conn: with db.conn:
db[table].insert_all( 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: except UnicodeDecodeErrorForPath as e:
raise click.ClickException( raise click.ClickException(

12
tests/test_insert_files.py
View file

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