New upsert implementation, refs #66

sqlite_utils/cli.py

@@ -378,6 +378,8 @@ def insert_upsert_implementation(
         extra_kwargs["not_null"] = set(not_null)
     if default:
         extra_kwargs["defaults"] = dict(default)
+    if upsert:
+        extra_kwargs["upsert"] = upsert
     db[table].insert_all(
         docs, pk=pk, batch_size=batch_size, alter=alter, **extra_kwargs
     )

sqlite_utils/db.py

@@ -949,6 +949,7 @@ class Table(Queryable):
         ignore=DEFAULT,
         replace=DEFAULT,
         extracts=DEFAULT,
+        upsert=False,
     ):
         """
         Like .insert() but takes a list of records and ensures that the table
@@ -985,8 +986,6 @@ class Table(Queryable):
         assert (
             num_columns <= SQLITE_MAX_VARS
         ), "Rows can have a maximum of {} columns".format(SQLITE_MAX_VARS)
-        # When calculating this for upsert, the primary keys are referenced twice
-        # so num_columns needs to have num primary keys added to it:
         batch_size = max(1, min(batch_size, SQLITE_MAX_VARS // num_columns))
         for chunk in chunks(itertools.chain([first_record], records), batch_size):
             chunk = list(chunk)
@@ -1011,37 +1010,9 @@ class Table(Queryable):
                     all_columns.insert(0, hash_id)
             first = False
 
-            # START of bit that differs for upsert v.s. insert
-            or_what = ""
-            if replace:
-                or_what = "OR REPLACE "
-            elif ignore:
-                or_what = "OR IGNORE "
-            # INSERT OR IGNORE INTO book(id) VALUES(1001);
-            # UPDATE book SET name = 'Programming' WHERE id = 1001;
-            # WAIT NO: this won't work because here we create a single
-            # giant INSERT, but for upsert we need two statements per
-            # record which means we need executescript()... but
-            # executescript doesn't take parameters at all.
-            # So maybe for upsert() we execute two separate SQL queries
-            # for every single record? Very different implementation.
-            sql = """
-                INSERT {or_what}INTO [{table}] ({columns}) VALUES {rows};
-            """.format(
-                or_what=or_what,
-                table=self.name,
-                columns=", ".join("[{}]".format(c) for c in all_columns),
-                rows=", ".join(
-                    """
-                    ({placeholders})
-                """.format(
-                        placeholders=", ".join(["?"] * len(all_columns))
-                    )
-                    for record in chunk
-                ),
-            )
-            # END of bit that differs
-
+            # values is the list of insert data that is passed to the
+            # .execute() method - but some of them may be replaced by
+            # new primary keys if we are extracting any columns.
             values = []
             extracts = resolve_extracts(extracts)
             for record in chunk:
@@ -1054,25 +1025,76 @@ class Table(Queryable):
                         extract_table = extracts[key]
                         value = self.db[extract_table].lookup({"value": value})
                     record_values.append(value)
-                values.extend(record_values)
+                values.append(record_values)
+
+            queries_and_params = []
+            if upsert:
+                if isinstance(pk, str):
+                    pks = [pk]
+                else:
+                    pks = pk
+                for record_values in values:
+                    # TODO: make more efficient:
+                    record = dict(zip(all_columns, record_values))
+                    params = []
+                    sql = "INSERT OR IGNORE INTO [{table}]({pks}) VALUES({pk_placeholders});".format(
+                        table=self.name,
+                        pks=", ".join(["[{}]".format(p) for p in pks]),
+                        pk_placeholders=", ".join(["?" for p in pks]),
+                    )
+                    queries_and_params.append((sql, [record[col] for col in pks]))
+                    # UPDATE [book] SET [name] = 'Programming' WHERE [id] = 1001;
+                    set_cols = [col for col in all_columns if col not in pks]
+                    sql2 = "UPDATE [{table}] SET {pairs} WHERE {wheres}".format(
+                        table=self.name,
+                        pairs=", ".join("[{}] = ?".format(col) for col in set_cols),
+                        wheres=" AND ".join("[{}] = ?".format(pk) for pk in pks),
+                    )
+                    queries_and_params.append(
+                        (
+                            sql2,
+                            [record[col] for col in set_cols]
+                            + [record[pk] for pk in pks],
+                        )
+                    )
+            else:
+                or_what = ""
+                if replace:
+                    or_what = "OR REPLACE "
+                elif ignore:
+                    or_what = "OR IGNORE "
+                sql = """
+                    INSERT {or_what}INTO [{table}] ({columns}) VALUES {rows};
+                """.format(
+                    or_what=or_what,
+                    table=self.name,
+                    columns=", ".join("[{}]".format(c) for c in all_columns),
+                    rows=", ".join(
+                        """
+                        ({placeholders})
+                    """.format(
+                            placeholders=", ".join(["?"] * len(all_columns))
+                        )
+                        for record in chunk
+                    ),
+                )
+                flat_values = list(itertools.chain(*values))
+                queries_and_params = [(sql, flat_values)]
 
-            # Except... if upsert() needs to execute multiple queries
-            # then this bit needs to change too
             with self.db.conn:
-                try:
-                    result = self.db.conn.execute(sql, values)
-                except OperationalError as e:
-                    if alter and (" column" in e.args[0]):
-                        # Attempt to add any missing columns, then try again
-                        self.add_missing_columns(chunk)
-                        result = self.db.conn.execute(sql, values)
-                    else:
-                        raise
-                # How do we get lastrowid for upserts?
+                for query, params in queries_and_params:
+                    try:
+                        result = self.db.conn.execute(query, params)
+                    except OperationalError as e:
+                        if alter and (" column" in e.args[0]):
+                            # Attempt to add any missing columns, then try again
+                            self.add_missing_columns(chunk)
+                            result = self.db.conn.execute(query, params)
+                        else:
+                            raise
                 self.last_rowid = result.lastrowid
                 self.last_pk = self.last_rowid
                 # self.last_rowid will be 0 if a "INSERT OR IGNORE" happened
-                # Do we need this logic for upsert too?
                 if (hash_id or pk) and self.last_rowid:
                     row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0]
                     if hash_id:
@@ -1123,7 +1145,19 @@ class Table(Queryable):
         # Perform the following for each record:
         # INSERT OR IGNORE INTO books(id) VALUES(1001);
         # UPDATE books SET name = 'Programming' WHERE id = 1001;
-        raise NotImplementedError
+        return self.insert_all(
+            records,
+            pk=pk,
+            foreign_keys=foreign_keys,
+            column_order=column_order,
+            not_null=not_null,
+            defaults=defaults,
+            batch_size=batch_size,
+            hash_id=hash_id,
+            alter=alter,
+            extracts=extracts,
+            upsert=True,
+        )
 
     def add_missing_columns(self, records):
         needed_columns = self.detect_column_types(records)

tests/test_cli.py

@@ -765,3 +765,59 @@ def test_rows(db_path, args, expected):
         )
     result = CliRunner().invoke(cli.cli, ["rows", db_path, "dogs"] + args)
     assert expected == result.output.strip()
+
+
+def test_upsert(db_path, tmpdir):
+    json_path = str(tmpdir / "dogs.json")
+    db = Database(db_path)
+    insert_dogs = [
+        {"id": 1, "name": "Cleo", "age": 4},
+        {"id": 2, "name": "Nixie", "age": 4},
+    ]
+    open(json_path, "w").write(json.dumps(insert_dogs))
+    result = CliRunner().invoke(
+        cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id"]
+    )
+    assert 0 == result.exit_code, result.output
+    assert 2 == db["dogs"].count
+    # Now run the upsert to update just their ages
+    upsert_dogs = [
+        {"id": 1, "age": 5},
+        {"id": 2, "age": 5},
+    ]
+    open(json_path, "w").write(json.dumps(insert_dogs))
+    result = CliRunner().invoke(
+        cli.cli, ["upsert", db_path, "dogs", json_path, "--pk", "id"]
+    )
+    assert 0 == result.exit_code, result.output
+    assert [
+        {"id": 1, "name": "Cleo", "age": 4},
+        {"id": 2, "name": "Nixie", "age": 4},
+    ] == db.execute_returning_dicts("select * from dogs order by id")
+
+
+def test_upsert_alter(db_path, tmpdir):
+    json_path = str(tmpdir / "dogs.json")
+    db = Database(db_path)
+    insert_dogs = [{"id": 1, "name": "Cleo"}]
+    open(json_path, "w").write(json.dumps(insert_dogs))
+    result = CliRunner().invoke(
+        cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id"]
+    )
+    assert 0 == result.exit_code, result.output
+    # Should fail with error code if no --alter
+    upsert_dogs = [{"id": 1, "age": 5}]
+    open(json_path, "w").write(json.dumps(upsert_dogs))
+    result = CliRunner().invoke(
+        cli.cli, ["upsert", db_path, "dogs", json_path, "--pk", "id"]
+    )
+    assert 1 == result.exit_code
+    assert "no such column: age" == str(result.exception)
+    # Should succeed with --alter
+    result = CliRunner().invoke(
+        cli.cli, ["upsert", db_path, "dogs", json_path, "--pk", "id", "--alter"]
+    )
+    assert 0 == result.exit_code
+    assert [{"id": 1, "name": "Cleo", "age": 5},] == db.execute_returning_dicts(
+        "select * from dogs order by id"
+    )

tests/test_upsert.py

@@ -0,0 +1,16 @@
+def test_upsert(fresh_db):
+    table = fresh_db["table"]
+    table.insert({"id": 1, "name": "Cleo"}, pk="id")
+    table.upsert({"id": 1, "age": 5}, pk="id", alter=True)
+    assert [{"id": 1, "name": "Cleo", "age": 5}] == list(table.rows)
+
+
+def test_upsert_all(fresh_db):
+    table = fresh_db["table"]
+    table.upsert_all([{"id": 1, "name": "Cleo"}, {"id": 2, "name": "Nixie"}], pk="id")
+    table.upsert_all([{"id": 1, "age": 5}, {"id": 2, "age": 5}], pk="id", alter=True)
+    assert [
+        {"id": 1, "name": "Cleo", "age": 5},
+        {"id": 2, "name": "Nixie", "age": 5},
+    ] == list(table.rows)
+    assert 2 == table.last_pk