diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index bc2e579..3e42480 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -62,24 +62,31 @@ def optimize(path, no_vacuum): db.vacuum() -@cli.command() -@click.argument( - "path", - type=click.Path(file_okay=True, dir_okay=False, allow_dash=False), - required=True, -) -@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("--nl", is_flag=True, help="Expect newline-delimited JSON") -@click.option("--csv", is_flag=True, help="Expect CSV") -@click.option("--batch-size", type=int, default=100, help="Commit every X records") -def insert(path, table, json_file, pk, nl, csv, batch_size): - """ - Insert records from JSON file into a table, creating the table if it does not already exist. +def insert_upsert_options(fn): + for decorator in reversed( + ( + click.argument( + "path", + type=click.Path(file_okay=True, dir_okay=False, allow_dash=False), + required=True, + ), + 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("--nl", is_flag=True, help="Expect newline-delimited JSON"), + click.option("--csv", is_flag=True, help="Expect CSV"), + click.option( + "--batch-size", type=int, default=100, help="Commit every X records" + ), + ) + ): + fn = decorator(fn) + return fn - Input should be a JSON array of objects, unless --nl or --csv is used. - """ + +def insert_upsert_implementation( + path, table, json_file, pk, nl, csv, batch_size, upsert +): db = sqlite_utils.Database(path) if nl and csv: click.echo("Use just one of --nl and --csv", err=True) @@ -94,25 +101,38 @@ def insert(path, table, json_file, pk, nl, csv, batch_size): docs = json.load(json_file) if isinstance(docs, dict): docs = [docs] - db[table].insert_all(docs, pk=pk, batch_size=batch_size) + if upsert: + method = db[table].upsert_all + else: + method = db[table].insert_all + method(docs, pk=pk, batch_size=batch_size) @cli.command() -@click.argument( - "path", - type=click.Path(file_okay=True, dir_okay=False, allow_dash=False), - required=True, -) -@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") -def upsert(path, table, json_file, pk): - "Upsert records based on their primary key" - db = sqlite_utils.Database(path) - docs = json.load(json_file) - if isinstance(docs, dict): - docs = [docs] - db[table].upsert_all(docs, pk=pk) +@insert_upsert_options +def insert(path, table, json_file, pk, nl, csv, batch_size): + """ + Insert records from JSON file into a table, creating the table if it + does not already exist. + + Input should be a JSON array of objects, unless --nl or --csv is used. + """ + insert_upsert_implementation( + path, table, json_file, pk, nl, csv, batch_size, upsert=False + ) + + +@cli.command() +@insert_upsert_options +def upsert(path, table, json_file, pk, nl, csv, batch_size): + """ + Upsert records based on their primary key. Works like 'insert' but if + an incoming record has a primary key that matches an existing record + the existing record will be replaced. + """ + insert_upsert_implementation( + path, table, json_file, pk, nl, csv, batch_size, upsert=True + ) @cli.command(name="csv") diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 963b1ae..6c8fa53 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -401,13 +401,16 @@ class Table: column_order=column_order, ) - def upsert_all(self, records, pk=None, foreign_keys=None, column_order=None): + def upsert_all( + self, records, pk=None, foreign_keys=None, column_order=None, batch_size=100 + ): return self.insert_all( records, pk=pk, foreign_keys=foreign_keys, - upsert=True, column_order=column_order, + batch_size=100, + upsert=True, )