WIP new upsert implementation, refs #652

docs/python-api.rst

@@ -927,6 +927,13 @@ An ``upsert_all()`` method is also available, which behaves like ``insert_all()`
 .. note::
     ``.upsert()`` and ``.upsert_all()`` in sqlite-utils 1.x worked like ``.insert(..., replace=True)`` and ``.insert_all(..., replace=True)`` do in 2.x. See `issue #66 <https://github.com/simonw/sqlite-utils/issues/66>`__ for details of this change.
 
+.. _python_api_old_upsert:
+
+Alternative upserts using INSERT OR IGNORE
+------------------------------------------
+
+Upserts use ``INSERT INTO ... ON CONFLICT SET``. Prior to ``sqlite-utils 3.x`` (TODO: fill in version) these used a sequence of ``INSERT OR IGNORE`` followed by an ``UPDATE``. This older method is still used for SQLite 3.23.1 and earlier. You can force the older implementation by passing ``use_old_upsert=True`` to the ``Database()`` constructor.
+
 .. _python_api_convert:
 
 Converting data in columns

sqlite_utils/db.py

@@ -304,6 +304,8 @@ class Database:
       ``sql, parameters`` every time a SQL query is executed
     :param use_counts_table: set to ``True`` to use a cached counts table, if available. See
       :ref:`python_api_cached_table_counts`
+    :param use_old_upsert: set to ``True`` to force the older upsert implementation. See
+      :ref:`python_api_old_upsert`
     :param strict: Apply STRICT mode to all created tables (unless overridden)
     """
 
@@ -320,10 +322,12 @@ class Database:
         tracer: Optional[Callable] = None,
         use_counts_table: bool = False,
         execute_plugins: bool = True,
+        use_old_upsert: bool = False,
         strict: bool = False,
     ):
         self.memory_name = None
         self.memory = False
+        self.use_old_upsert = use_old_upsert
         assert (filename_or_conn is not None and (not memory and not memory_name)) or (
             filename_or_conn is None and (memory or memory_name)
         ), "Either specify a filename_or_conn or pass memory=True"
@@ -684,6 +688,33 @@ class Database:
                 self._supports_strict = False
         return self._supports_strict
 
+    @property
+    def supports_on_conflict(self) -> bool:
+        # SQLite's upsert is implemented as INSERT INTO ... ON CONFLICT DO ...
+        if not hasattr(self, "_supports_on_conflict"):
+            try:
+                table_name = "t{}".format(secrets.token_hex(16))
+                with self.conn:
+                    self.conn.execute(
+                        "create table {} (id integer primary key, name text)".format(
+                            table_name
+                        )
+                    )
+                    self.conn.execute(
+                        "insert into {} (id, name) values (1, 'one')".format(table_name)
+                    )
+                    self.conn.execute(
+                        (
+                            "insert into {} (id, name) values (1, 'two') "
+                            "on conflict do update set name = 'two'"
+                        ).format(table_name)
+                    )
+                    self.conn.execute("drop table {}".format(table_name))
+                self._supports_on_conflict = True
+            except Exception:
+                self._supports_on_conflict = False
+        return self._supports_on_conflict
+
     @property
     def sqlite_version(self) -> Tuple[int, ...]:
         "Version of SQLite, as a tuple of integers for example ``(3, 36, 0)``."
@@ -2968,102 +2999,128 @@ class Table(Queryable):
         replace,
         ignore,
     ):
-        # values is the list of insert data that is passed to the
+        """
-        # .execute() method - but some of them may be replaced by
+        Given a list ``chunk`` of records that should be written to *this* table,
-        # new primary keys if we are extracting any columns.
+        return a list of ``(sql, parameters)`` 2-tuples which, when executed in
-        values = []
+        order, perform the desired INSERT / UPSERT / REPLACE operation.
+        """
+        if (upsert or replace) and not pk:
+            raise PrimaryKeyRequired("UPSERT / REPLACE needs a primary key")
+
         if hash_id_columns and hash_id is None:
             hash_id = "id"
+
         extracts = resolve_extracts(extracts)
+
+        # Build a row-list ready for executemany-style flattening
+        values: list[list] = []
         for record in chunk:
-            record_values = []
+            row_vals = []
-            for key in all_columns:
+            for col in all_columns:
-                value = jsonify_if_needed(
+                if col == hash_id:
-                    record.get(
+                    row_vals.append(hash_record(record, hash_id_columns))
-                        key,
+                    continue
-                        (
-                            None
-                            if key != hash_id
-                            else hash_record(record, hash_id_columns)
-                        ),
-                    )
-                )
-                if key in extracts:
-                    extract_table = extracts[key]
-                    value = self.db[extract_table].lookup({"value": value})
-                record_values.append(value)
-            values.append(record_values)
 
-        queries_and_params = []
+                val = record.get(col)
-        if upsert:
+                if val is None and not_null and col in not_null:
-            if isinstance(pk, str):
+                    val = ""
-                pks = [pk]
+                row_vals.append(jsonify_if_needed(val))
-            else:
+            values.append(row_vals)
-                pks = pk
-            self.last_pk = None
-            for record_values in values:
-                record = dict(zip(all_columns, record_values))
-                placeholders = list(pks)
-                # Need to populate not-null columns too, or INSERT OR IGNORE ignores
-                # them since it ignores the resulting integrity errors
-                if not_null:
-                    placeholders.extend(not_null)
-                sql = "INSERT OR IGNORE INTO [{table}]({cols}) VALUES({placeholders});".format(
-                    table=self.name,
-                    cols=", ".join(["[{}]".format(p) for p in placeholders]),
-                    placeholders=", ".join(["?" for p in placeholders]),
-                )
-                queries_and_params.append(
-                    (sql, [record[col] for col in pks] + ["" for _ in (not_null or [])])
-                )
-                # UPDATE [book] SET [name] = 'Programming' WHERE [id] = 1001;
-                set_cols = [col for col in all_columns if col not in pks]
-                if set_cols:
-                    sql2 = "UPDATE [{table}] SET {pairs} WHERE {wheres}".format(
-                        table=self.name,
-                        pairs=", ".join(
-                            "[{}] = {}".format(col, conversions.get(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],
-                        )
-                    )
-                # We can populate .last_pk right here
-                if num_records_processed == 1:
-                    self.last_pk = tuple(record[pk] for pk in pks)
-                    if len(self.last_pk) == 1:
-                        self.last_pk = self.last_pk[0]
 
-        else:
+        columns_sql = ", ".join(f"[{c}]" for c in all_columns)
-            or_what = ""
+        placeholder_expr = ", ".join(conversions.get(c, "?") for c in all_columns)
-            if replace:
+        row_placeholders_sql = ", ".join(f"({placeholder_expr})" for _ in values)
-                or_what = "OR REPLACE "
+        flat_params = list(itertools.chain.from_iterable(values))
-            elif ignore:
+
-                or_what = "OR IGNORE "
+        # replace=True mean INSERT OR REPLACE INTO
-            sql = """
+        if replace:
-                INSERT {or_what}INTO [{table}] ({columns}) VALUES {rows};
+            sql = (
-            """.strip().format(
+                f"INSERT OR REPLACE INTO [{self.name}] "
-                or_what=or_what,
+                f"({columns_sql}) VALUES {row_placeholders_sql}"
-                table=self.name,
-                columns=", ".join("[{}]".format(c) for c in all_columns),
-                rows=", ".join(
-                    "({placeholders})".format(
-                        placeholders=", ".join(
-                            [conversions.get(col, "?") for col in all_columns]
-                        )
-                    )
-                    for record in chunk
-                ),
             )
-            flat_values = list(itertools.chain(*values))
+            return [(sql, flat_params)]
-            queries_and_params = [(sql, flat_values)]
 
+        # If not an upsert it's an INSERT, maybe with OR IGNORE
+        if not upsert:
+            or_ignore = ""
+            if ignore:
+                or_ignore = " OR IGNORE"
+            sql = (
+                f"INSERT{or_ignore} INTO [{self.name}] "
+                f"({columns_sql}) VALUES {row_placeholders_sql}"
+            )
+            return [(sql, flat_params)]
+
+        # Everything from here on is for upsert=True
+        pk_cols = [pk] if isinstance(pk, str) else list(pk)
+        non_pk_cols = [c for c in all_columns if c not in pk_cols]
+        conflict_sql = ", ".join(f"[{c}]" for c in pk_cols)
+
+        if self.db.supports_on_conflict and not self.db.use_old_upsert:
+            if non_pk_cols:
+                # DO UPDATE
+                assignments = []
+                for c in non_pk_cols:
+                    if c in conversions:
+                        assignments.append(
+                            f"[{c}] = {conversions[c].replace('?', f'excluded.[{c}]')}"
+                        )
+                    else:
+                        assignments.append(f"[{c}] = excluded.[{c}]")
+                do_clause = "DO UPDATE SET " + ", ".join(assignments)
+            else:
+                # All columns are in the PK – nothing to update.
+                do_clause = "DO NOTHING"
+
+            sql = (
+                f"INSERT INTO [{self.name}] ({columns_sql}) "
+                f"VALUES {row_placeholders_sql} "
+                f"ON CONFLICT({conflict_sql}) {do_clause}"
+            )
+            return [(sql, flat_params)]
+
+        # At this point we need compatibility UPSERT for SQLite < 3.24.0
+        # (INSERT OR IGNORE + second UPDATE stage)
+        queries_and_params: list[tuple[str, list]] = []
+
+        insert_sql = (
+            f"INSERT OR IGNORE INTO [{self.name}] "
+            f"({columns_sql}) VALUES {row_placeholders_sql}"
+        )
+        queries_and_params.append((insert_sql, flat_params))
+
+        # If there is nothing to update we are done.
+        if not non_pk_cols:
+            return queries_and_params
+
+        # We can use UPDATE … FROM (VALUES …) on SQLite ≥ 3.33.0
+        # Older SQLite versions will run this as one UPDATE per row
+        # – which is what sqlite-utils did prior to this refactor.
+        alias_cols_sql = ", ".join(pk_cols + non_pk_cols)
+
+        assignments = []
+        for c in non_pk_cols:
+            if c in conversions:
+                assignments.append(f"[{c}] = {conversions[c].replace('?', f'v.[{c}]')}")
+            else:
+                assignments.append(f"[{c}] = v.[{c}]")
+        assignments_sql = ", ".join(assignments)
+
+        update_sql = (
+            f"UPDATE [{self.name}] AS m SET {assignments_sql} "
+            f"FROM (VALUES {row_placeholders_sql}) "
+            f"AS v({alias_cols_sql}) "
+            f"WHERE " + " AND ".join(f"m.[{c}] = v.[{c}]" for c in pk_cols)
+        )
+
+        # Parameters for the UPDATE – pk cols first then non-pk cols
+        update_params: list = []
+        for row in values:
+            row_dict = dict(zip(all_columns, row))
+            ordered = [row_dict[c] for c in pk_cols + non_pk_cols]
+            update_params.extend(ordered)
+
+        queries_and_params.append((update_sql, update_params))
         return queries_and_params
 
     def insert_chunk(
@@ -3081,7 +3138,7 @@ class Table(Queryable):
         num_records_processed,
         replace,
         ignore,
-    ):
+    ) -> Optional[sqlite3.Cursor]:
         queries_and_params = self.build_insert_queries_and_params(
             extracts,
             chunk,
@@ -3096,9 +3153,8 @@ class Table(Queryable):
             replace,
             ignore,
         )
-
+        result = None
         with self.db.conn:
-            result = None
             for query, params in queries_and_params:
                 try:
                     result = self.db.execute(query, params)
@@ -3127,7 +3183,7 @@ class Table(Queryable):
                             ignore,
                         )
 
-                        self.insert_chunk(
+                        result = self.insert_chunk(
                             alter,
                             extracts,
                             second_half,
@@ -3145,20 +3201,7 @@ class Table(Queryable):
 
                     else:
                         raise
-            if num_records_processed == 1 and not upsert:
+        return result
-                self.last_rowid = result.lastrowid
-                self.last_pk = self.last_rowid
-                # self.last_rowid will be 0 if a "INSERT OR IGNORE" happened
-                if (hash_id or pk) and self.last_rowid:
-                    row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0]
-                    if hash_id:
-                        self.last_pk = row[hash_id]
-                    elif isinstance(pk, str):
-                        self.last_pk = row[pk]
-                    else:
-                        self.last_pk = tuple(row[p] for p in pk)
-
-        return
 
     def insert(
         self,
@@ -3309,6 +3352,7 @@ class Table(Queryable):
         self.last_pk = None
         if truncate and self.exists():
             self.db.execute("DELETE FROM [{}];".format(self.name))
+        result = None
         for chunk in chunks(itertools.chain([first_record], records), batch_size):
             chunk = list(chunk)
             num_records_processed += len(chunk)
@@ -3316,6 +3360,12 @@ class Table(Queryable):
                 if not self.exists():
                     # Use the first batch to derive the table names
                     column_types = suggest_column_types(chunk)
+                    if extracts:
+                        for col in extracts:
+                            if col in column_types:
+                                column_types[col] = (
+                                    int  # This will be an integer foreign key
+                                )
                     column_types.update(columns or {})
                     self.create(
                         column_types,
@@ -3343,7 +3393,7 @@ class Table(Queryable):
 
             first = False
 
-            self.insert_chunk(
+            result = self.insert_chunk(
                 alter,
                 extracts,
                 chunk,
@@ -3359,6 +3409,33 @@ class Table(Queryable):
                 ignore,
             )
 
+        # If we only handled a single row populate self.last_pk
+        if num_records_processed == 1:
+            # For an insert we need to use result.lastrowid
+            if not upsert:
+                self.last_rowid = result.lastrowid
+                if (hash_id or pk) and self.last_rowid:
+                    # Set self.last_pk to the pk(s) for that rowid
+                    row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0]
+                    if hash_id:
+                        self.last_pk = row[hash_id]
+                    elif isinstance(pk, str):
+                        self.last_pk = row[pk]
+                    else:
+                        self.last_pk = tuple(row[p] for p in pk)
+                else:
+                    self.last_pk = self.last_rowid
+            else:
+                # For an upsert use first_record from earlier
+                if hash_id:
+                    self.last_pk = hash_record(first_record, hash_id_columns)
+                else:
+                    self.last_pk = (
+                        first_record[pk]
+                        if isinstance(pk, str)
+                        else tuple(first_record[p] for p in pk)
+                    )
+
         if analyze:
             self.analyze()
 

tests/test_cli.py

@@ -1117,9 +1117,8 @@ def test_upsert_alter(db_path, tmpdir):
     )
     assert result.exit_code == 1
     assert (
-        "Error: no such column: age\n\n"
+        "Error: table dogs has no column named age\n\n"
-        "sql = UPDATE [dogs] SET [age] = ? WHERE [id] = ?\n"
+        "Try using --alter to add additional columns"
-        "parameters = [5, 1]"
     ) == result.output.strip()
     # Should succeed with --alter
     result = CliRunner().invoke(
@@ -2248,7 +2247,7 @@ def test_integer_overflow_error(tmpdir):
     assert result.exit_code == 1
     assert result.output == (
         "Error: Python int too large to convert to SQLite INTEGER\n\n"
-        "sql = INSERT INTO [items] ([bignumber]) VALUES (?);\n"
+        "sql = INSERT INTO [items] ([bignumber]) VALUES (?)\n"
         "parameters = [34223049823094832094802398430298048240]\n"
     )
 

tests/test_tracer.py

@@ -18,7 +18,7 @@ def test_tracer():
         ("select name from sqlite_master where type = 'view'", None),
         ("CREATE TABLE [dogs] (\n   [name] TEXT\n);\n        ", None),
         ("select name from sqlite_master where type = 'view'", None),
-        ("INSERT INTO [dogs] ([name]) VALUES (?);", ["Cleopaws"]),
+        ("INSERT INTO [dogs] ([name]) VALUES (?)", ["Cleopaws"]),
         ("select name from sqlite_master where type = 'view'", None),
         (
             "CREATE VIRTUAL TABLE [dogs_fts] USING FTS5 (\n    [name],\n    content=[dogs]\n)",