From fe3a3d4000c9825f323b83276a1f9b3396dbbe3b Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Wed, 7 May 2025 17:10:22 -0700 Subject: [PATCH] WIP new upsert implementation, refs #652 --- docs/python-api.rst | 7 ++ sqlite_utils/db.py | 291 +++++++++++++++++++++++++++---------------- tests/test_cli.py | 7 +- tests/test_tracer.py | 2 +- 4 files changed, 195 insertions(+), 112 deletions(-) diff --git a/docs/python-api.rst b/docs/python-api.rst index c6bf776..6893e1c 100644 --- a/docs/python-api.rst +++ b/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 `__ 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 diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 927e6c6..b960790 100644 --- a/sqlite_utils/db.py +++ b/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 - # new primary keys if we are extracting any columns. - values = [] + """ + Given a list ``chunk`` of records that should be written to *this* table, + return a list of ``(sql, parameters)`` 2-tuples which, when executed in + 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 = [] - for key in all_columns: - value = jsonify_if_needed( - record.get( - key, - ( - 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) + row_vals = [] + for col in all_columns: + if col == hash_id: + row_vals.append(hash_record(record, hash_id_columns)) + continue - queries_and_params = [] - if upsert: - if isinstance(pk, str): - pks = [pk] - else: - 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] + val = record.get(col) + if val is None and not_null and col in not_null: + val = "" + row_vals.append(jsonify_if_needed(val)) + values.append(row_vals) - else: - or_what = "" - if replace: - or_what = "OR REPLACE " - elif ignore: - or_what = "OR IGNORE " - sql = """ - INSERT {or_what}INTO [{table}] ({columns}) VALUES {rows}; - """.strip().format( - or_what=or_what, - 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 - ), + columns_sql = ", ".join(f"[{c}]" for c in all_columns) + placeholder_expr = ", ".join(conversions.get(c, "?") for c in all_columns) + row_placeholders_sql = ", ".join(f"({placeholder_expr})" for _ in values) + flat_params = list(itertools.chain.from_iterable(values)) + + # replace=True mean INSERT OR REPLACE INTO + if replace: + sql = ( + f"INSERT OR REPLACE INTO [{self.name}] " + f"({columns_sql}) VALUES {row_placeholders_sql}" ) - flat_values = list(itertools.chain(*values)) - queries_and_params = [(sql, flat_values)] + return [(sql, flat_params)] + # 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: - 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 + return result 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() diff --git a/tests/test_cli.py b/tests/test_cli.py index 4033af6..88763ee 100644 --- a/tests/test_cli.py +++ b/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" - "sql = UPDATE [dogs] SET [age] = ? WHERE [id] = ?\n" - "parameters = [5, 1]" + "Error: table dogs has no column named age\n\n" + "Try using --alter to add additional columns" ) == 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" ) diff --git a/tests/test_tracer.py b/tests/test_tracer.py index 26318ae..9dfb490 100644 --- a/tests/test_tracer.py +++ b/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)",