From bf354b776cc664edf43208079b09dd2c23beeab7 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 23 Nov 2025 09:16:33 -0800 Subject: [PATCH] Initial Claude implementation of list mode From https://github.com/simonw/research/tree/main/sqlite-utils-iterator-support I made very minor changes to the test (removing a sys.path insert) Refs #672 Refs https://github.com/simonw/research/pull/31 --- sqlite_utils/db.py | 138 ++++++++++++++++++++++--------- tests/test_list_mode.py | 175 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 276 insertions(+), 37 deletions(-) create mode 100644 tests/test_list_mode.py diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 2be7a6d..02962ac 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -3010,6 +3010,7 @@ class Table(Queryable): num_records_processed, replace, ignore, + list_mode=False, ): """ Given a list ``chunk`` of records that should be written to *this* table, @@ -3024,24 +3025,40 @@ class Table(Queryable): # Build a row-list ready for executemany-style flattening values = [] - 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 list_mode: + # In list mode, records are already lists of values + for record in chunk: + record_values = [] + for i, key in enumerate(all_columns): + if i < len(record): + value = jsonify_if_needed(record[i]) + else: + value = None + if key in extracts: + extract_table = extracts[key] + value = self.db[extract_table].lookup({"value": value}) + record_values.append(value) + values.append(record_values) + else: + # Dict mode: original logic + 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) + if key in extracts: + extract_table = extracts[key] + value = self.db[extract_table].lookup({"value": value}) + record_values.append(value) + values.append(record_values) columns_sql = ", ".join(f"[{c}]" for c in all_columns) placeholder_expr = ", ".join(conversions.get(c, "?") for c in all_columns) @@ -3157,6 +3174,7 @@ class Table(Queryable): num_records_processed, replace, ignore, + list_mode=False, ) -> Optional[sqlite3.Cursor]: queries_and_params = self.build_insert_queries_and_params( extracts, @@ -3171,6 +3189,7 @@ class Table(Queryable): num_records_processed, replace, ignore, + list_mode, ) result = None with self.db.conn: @@ -3200,6 +3219,7 @@ class Table(Queryable): num_records_processed, replace, ignore, + list_mode, ) result = self.insert_chunk( @@ -3216,6 +3236,7 @@ class Table(Queryable): num_records_processed, replace, ignore, + list_mode, ) else: @@ -3353,17 +3374,47 @@ class Table(Queryable): all_columns = [] first = True num_records_processed = 0 - # Fix up any records with square braces in the column names - records = fix_square_braces(records) - # We can only handle a max of 999 variables in a SQL insert, so - # we need to adjust the batch_size down if we have too many cols - records = iter(records) - # Peek at first record to count its columns: + + # Detect if we're using list-based iteration or dict-based iteration + list_mode = False + column_names = None + + # Fix up any records with square braces in the column names (only for dict mode) + # We'll handle this differently for list mode + records_iter = iter(records) + + # Peek at first record to determine mode: try: - first_record = next(records) + first_record = next(records_iter) except StopIteration: return self # It was an empty list - num_columns = len(first_record.keys()) + + # Check if this is list mode or dict mode + if isinstance(first_record, list): + # List mode: first record should be column names + list_mode = True + if not all(isinstance(col, str) for col in first_record): + raise ValueError("When using list-based iteration, the first yielded value must be a list of column name strings") + column_names = first_record + all_columns = column_names + num_columns = len(column_names) + # Get the actual first data record + try: + first_record = next(records_iter) + except StopIteration: + return self # Only headers, no data + if not isinstance(first_record, list): + raise ValueError("After column names list, all subsequent records must also be lists") + else: + # Dict mode: traditional behavior + records_iter = itertools.chain([first_record], records_iter) + records_iter = fix_square_braces(records_iter) + try: + first_record = next(records_iter) + except StopIteration: + return self + num_columns = len(first_record.keys()) + assert ( num_columns <= SQLITE_MAX_VARS ), "Rows can have a maximum of {} columns".format(SQLITE_MAX_VARS) @@ -3373,13 +3424,18 @@ class Table(Queryable): 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): + for chunk in chunks(itertools.chain([first_record], records_iter), batch_size): chunk = list(chunk) num_records_processed += len(chunk) if first: if not self.exists(): # Use the first batch to derive the table names - column_types = suggest_column_types(chunk) + if list_mode: + # Convert list records to dicts for type detection + chunk_as_dicts = [dict(zip(column_names, row)) for row in chunk] + column_types = suggest_column_types(chunk_as_dicts) + else: + column_types = suggest_column_types(chunk) if extracts: for col in extracts: if col in column_types: @@ -3399,17 +3455,24 @@ class Table(Queryable): extracts=extracts, strict=strict, ) - all_columns_set = set() - for record in chunk: - all_columns_set.update(record.keys()) - all_columns = list(sorted(all_columns_set)) - if hash_id: - all_columns.insert(0, hash_id) + if list_mode: + # In list mode, columns are already known + all_columns = list(column_names) + if hash_id: + all_columns.insert(0, hash_id) + else: + all_columns_set = set() + for record in chunk: + all_columns_set.update(record.keys()) + all_columns = list(sorted(all_columns_set)) + if hash_id: + all_columns.insert(0, hash_id) else: - for record in chunk: - all_columns += [ - column for column in record if column not in all_columns - ] + if not list_mode: + for record in chunk: + all_columns += [ + column for column in record if column not in all_columns + ] first = False @@ -3427,6 +3490,7 @@ class Table(Queryable): num_records_processed, replace, ignore, + list_mode, ) # If we only handled a single row populate self.last_pk diff --git a/tests/test_list_mode.py b/tests/test_list_mode.py new file mode 100644 index 0000000..4f841ff --- /dev/null +++ b/tests/test_list_mode.py @@ -0,0 +1,175 @@ +""" +Tests for list-based iteration in insert_all and upsert_all +""" +import pytest +from sqlite_utils import Database + + +def test_insert_all_list_mode_basic(): + """Test basic insert_all with list-based iteration""" + db = Database(memory=True) + + def data_generator(): + # First yield column names + yield ["id", "name", "age"] + # Then yield data rows + yield [1, "Alice", 30] + yield [2, "Bob", 25] + yield [3, "Charlie", 35] + + db["people"].insert_all(data_generator()) + + rows = list(db["people"].rows) + assert len(rows) == 3 + assert rows[0] == {"id": 1, "name": "Alice", "age": 30} + assert rows[1] == {"id": 2, "name": "Bob", "age": 25} + assert rows[2] == {"id": 3, "name": "Charlie", "age": 35} + + +def test_insert_all_list_mode_with_pk(): + """Test insert_all with list mode and primary key""" + db = Database(memory=True) + + def data_generator(): + yield ["id", "name", "score"] + yield [1, "Alice", 95] + yield [2, "Bob", 87] + + db["scores"].insert_all(data_generator(), pk="id") + + assert db["scores"].pks == ["id"] + rows = list(db["scores"].rows) + assert len(rows) == 2 + + +def test_upsert_all_list_mode(): + """Test upsert_all with list-based iteration""" + db = Database(memory=True) + + # Initial insert + def initial_data(): + yield ["id", "name", "value"] + yield [1, "Alice", 100] + yield [2, "Bob", 200] + + db["data"].insert_all(initial_data(), pk="id") + + # Upsert with some updates and new records + def upsert_data(): + yield ["id", "name", "value"] + yield [1, "Alice", 150] # Update existing + yield [3, "Charlie", 300] # Insert new + + db["data"].upsert_all(upsert_data(), pk="id") + + rows = list(db["data"].rows_where(order_by="id")) + assert len(rows) == 3 + assert rows[0] == {"id": 1, "name": "Alice", "value": 150} + assert rows[1] == {"id": 2, "name": "Bob", "value": 200} + assert rows[2] == {"id": 3, "name": "Charlie", "value": 300} + + +def test_list_mode_with_various_types(): + """Test list mode with different data types""" + db = Database(memory=True) + + def data_generator(): + yield ["id", "name", "score", "active"] + yield [1, "Alice", 95.5, True] + yield [2, "Bob", 87.3, False] + yield [3, "Charlie", None, True] + + db["mixed"].insert_all(data_generator()) + + rows = list(db["mixed"].rows) + assert len(rows) == 3 + assert rows[0]["score"] == 95.5 + assert rows[1]["active"] == 0 # SQLite stores boolean as int + assert rows[2]["score"] is None + + +def test_list_mode_error_non_string_columns(): + """Test that non-string column names raise an error""" + db = Database(memory=True) + + def bad_data(): + yield [1, 2, 3] # Non-string column names + yield ["a", "b", "c"] + + with pytest.raises(ValueError, match="must be a list of column name strings"): + db["bad"].insert_all(bad_data()) + + +def test_list_mode_error_mixed_types(): + """Test that mixing list and dict raises an error""" + db = Database(memory=True) + + def bad_data(): + yield ["id", "name"] + yield {"id": 1, "name": "Alice"} # Should be a list, not dict + + with pytest.raises(ValueError, match="must also be lists"): + db["bad"].insert_all(bad_data()) + + +def test_list_mode_empty_after_headers(): + """Test that only headers without data works gracefully""" + db = Database(memory=True) + + def data_generator(): + yield ["id", "name", "age"] + # No data rows + + result = db["people"].insert_all(data_generator()) + assert result is not None + assert not db["people"].exists() + + +def test_list_mode_batch_processing(): + """Test list mode with large dataset requiring batching""" + db = Database(memory=True) + + def large_data(): + yield ["id", "value"] + for i in range(1000): + yield [i, f"value_{i}"] + + db["large"].insert_all(large_data(), batch_size=100) + + count = db.execute("SELECT COUNT(*) as c FROM large").fetchone()[0] + assert count == 1000 + + +def test_list_mode_shorter_rows(): + """Test that rows shorter than column list get NULL values""" + db = Database(memory=True) + + def data_generator(): + yield ["id", "name", "age", "city"] + yield [1, "Alice", 30, "NYC"] + yield [2, "Bob"] # Missing age and city + yield [3, "Charlie", 35] # Missing city + + db["people"].insert_all(data_generator()) + + rows = list(db["people"].rows_where(order_by="id")) + assert rows[0] == {"id": 1, "name": "Alice", "age": 30, "city": "NYC"} + assert rows[1] == {"id": 2, "name": "Bob", "age": None, "city": None} + assert rows[2] == {"id": 3, "name": "Charlie", "age": 35, "city": None} + + +def test_backwards_compatibility_dict_mode(): + """Ensure dict mode still works (backward compatibility)""" + db = Database(memory=True) + + # Traditional dict-based insert + data = [ + {"id": 1, "name": "Alice", "age": 30}, + {"id": 2, "name": "Bob", "age": 25}, + ] + + db["people"].insert_all(data) + + rows = list(db["people"].rows) + assert len(rows) == 2 + assert rows[0] == {"id": 1, "name": "Alice", "age": 30}