diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 9090a23..531d017 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -289,7 +289,7 @@ def insert_upsert_options(fn): def insert_upsert_implementation( - path, table, json_file, pk, nl, csv, batch_size, alter, upsert + path, table, json_file, pk, nl, csv, batch_size, alter, upsert, ignore=False ): db = sqlite_utils.Database(path) if nl and csv: @@ -307,14 +307,19 @@ def insert_upsert_implementation( docs = [docs] if upsert: method = db[table].upsert_all + extra_kwargs = {} else: method = db[table].insert_all - method(docs, pk=pk, batch_size=batch_size, alter=alter) + extra_kwargs = {"ignore": ignore} + method(docs, pk=pk, batch_size=batch_size, alter=alter, **extra_kwargs) @cli.command() @insert_upsert_options -def insert(path, table, json_file, pk, nl, csv, batch_size, alter): +@click.option( + "--ignore", is_flag=True, default=False, help="Ignore records if pk already exists" +) +def insert(path, table, json_file, pk, nl, csv, batch_size, alter, ignore): """ Insert records from JSON file into a table, creating the table if it does not already exist. @@ -322,7 +327,16 @@ def insert(path, table, json_file, pk, nl, csv, batch_size, alter): 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, alter=alter, upsert=False + path, + table, + json_file, + pk, + nl, + csv, + batch_size, + alter=alter, + upsert=False, + ignore=ignore, ) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 0e8ee6b..b354d22 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -446,6 +446,7 @@ class Table: column_order=None, hash_id=None, alter=False, + ignore=False, ): return self.insert_all( [record], @@ -455,6 +456,7 @@ class Table: column_order=column_order, hash_id=hash_id, alter=alter, + ignore=ignore, ) def insert_all( @@ -467,6 +469,7 @@ class Table: column_order=None, hash_id=None, alter=False, + ignore=False, ): """ Like .insert() but takes a list of records and ensures that the table @@ -474,6 +477,9 @@ class Table: data """ assert not (hash_id and pk), "Use either pk= or hash_id=" + assert not ( + ignore and upsert + ), "Use either ignore=True or upsert=True, not both" all_columns = None first = True for chunk in chunks(records, batch_size): @@ -495,10 +501,15 @@ class Table: if hash_id: all_columns.insert(0, hash_id) first = False + or_what = "" + if upsert: + or_what = "OR REPLACE " + elif ignore: + or_what = "OR IGNORE " sql = """ - INSERT {upsert} INTO [{table}] ({columns}) VALUES {rows}; + INSERT {or_what}INTO [{table}] ({columns}) VALUES {rows}; """.format( - upsert="OR REPLACE" if upsert else "", + or_what=or_what, table=self.name, columns=", ".join("[{}]".format(c) for c in all_columns), rows=", ".join( @@ -530,7 +541,8 @@ class Table: raise self.last_rowid = result.lastrowid self.last_pk = None - if hash_id or pk: + # self.last_rowid will be 0 if a "INSERT OR IGNORE" happened + if (hash_id or pk) and self.last_rowid: self.last_pk = self.db.conn.execute( "select [{}] from [{}] where rowid = ?".format( hash_id or pk, self.name diff --git a/tests/test_cli.py b/tests/test_cli.py index 9f4f801..ec30c0b 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -362,6 +362,27 @@ def test_insert_newline_delimited(db_path): ] == db.execute_returning_dicts("select foo, n from from_json_nl") +def test_insert_ignore(db_path, tmpdir): + db = Database(db_path) + db["dogs"].insert({"id": 1, "name": "Cleo"}, pk="id") + json_path = str(tmpdir / "dogs.json") + open(json_path, "w").write(json.dumps([{"id": 1, "name": "Bailey"}])) + # Should raise error without --ignore + result = CliRunner().invoke( + cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id"] + ) + assert 0 != result.exit_code, result.output + # If we use --ignore it should run OK + result = CliRunner().invoke( + cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id", "--ignore"] + ) + assert 0 == result.exit_code, result.output + # ... but it should actually have no effect + assert [{"id": 1, "name": "Cleo"}] == db.execute_returning_dicts( + "select * from dogs" + ) + + def test_upsert(db_path, tmpdir): test_insert_multiple_with_primary_key(db_path, tmpdir) json_path = str(tmpdir / "upsert.json") diff --git a/tests/test_create.py b/tests/test_create.py index 0420d7d..227d4dd 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -392,6 +392,18 @@ def test_insert_thousands_ignores_extra_columns_after_first_100(fresh_db): assert [{"i": 101, "word": None}] == rows +def test_insert_ignore(fresh_db): + fresh_db["test"].insert({"id": 1, "bar": 2}, pk="id") + # Should raise an error if we try this again + with pytest.raises(sqlite3.IntegrityError): + fresh_db["test"].insert({"id": 1, "bar": 2}, pk="id") + # Using ignore=True should cause our insert to be silently ignored + fresh_db["test"].insert({"id": 1, "bar": 3}, pk="id", ignore=True) + # Only one row, and it should be bar=2, not bar=3 + rows = fresh_db.execute_returning_dicts("select * from test") + assert [{"id": 1, "bar": 2}] == rows + + def test_insert_hash_id(fresh_db): dogs = fresh_db["dogs"] id = dogs.upsert({"name": "Cleo", "twitter": "cleopaws"}, hash_id="id").last_pk