diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index fe5cd643b..19295831f 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -423,6 +423,7 @@ def insert_upsert_implementation( upsert, ignore=False, replace=False, + truncate=False, not_null=None, default=None, ): @@ -442,7 +443,7 @@ def insert_upsert_implementation( docs = json.load(json_file) if isinstance(docs, dict): docs = [docs] - extra_kwargs = {"ignore": ignore, "replace": replace} + extra_kwargs = {"ignore": ignore, "replace": replace, "truncate": truncate} if not_null: extra_kwargs["not_null"] = set(not_null) if default: @@ -465,6 +466,12 @@ def insert_upsert_implementation( default=False, help="Replace records if pk already exists", ) +@click.option( + "--truncate", + is_flag=True, + default=False, + help="Truncate table before inserting records, if table already exists", +) def insert( path, table, @@ -477,6 +484,7 @@ def insert( alter, ignore, replace, + truncate, not_null, default, ): @@ -499,6 +507,7 @@ def insert( upsert=False, ignore=ignore, replace=replace, + truncate=truncate, not_null=not_null, default=default, ) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 5d8e3e8a4..4107ceb52 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -976,6 +976,7 @@ def insert_all( alter=DEFAULT, ignore=DEFAULT, replace=DEFAULT, + truncate=False, extracts=DEFAULT, conversions=DEFAULT, columns=DEFAULT, @@ -1032,6 +1033,8 @@ def insert_all( # issue implicit BEGINs for DDL, only DML. We mix DDL and DML # below and might execute DDL first, e.g. for table creation. self.db.conn.execute("BEGIN") + if truncate and self.exists(): + self.db.conn.execute("DELETE FROM [{}];".format(self.name)) for chunk in chunks(itertools.chain([first_record], records), batch_size): chunk = list(chunk) num_records_processed += len(chunk) diff --git a/tests/test_cli.py b/tests/test_cli.py index 9e43fb9ca..0b07f56b8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -642,6 +642,31 @@ def test_insert_replace(db_path, tmpdir): ) +def test_insert_truncate(db_path): + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "from_json_nl", "-", "--nl", "--batch-size=1"], + input='{"foo": "bar", "n": 1}\n{"foo": "baz", "n": 2}', + ) + assert 0 == result.exit_code, result.output + db = Database(db_path) + assert [ + {"foo": "bar", "n": 1}, + {"foo": "baz", "n": 2}, + ] == db.execute_returning_dicts("select foo, n from from_json_nl") + # Truncate and insert new rows + result = CliRunner().invoke( + cli.cli, + ["insert", db_path, "from_json_nl", "-", "--nl", "--truncate", "--batch-size=1"], + input='{"foo": "bam", "n": 3}\n{"foo": "bat", "n": 4}', + ) + assert 0 == result.exit_code, result.output + assert [ + {"foo": "bam", "n": 3}, + {"foo": "bat", "n": 4}, + ] == db.execute_returning_dicts("select foo, n from from_json_nl") + + def test_insert_alter(db_path, tmpdir): result = CliRunner().invoke( cli.cli,