mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-09-27 20:34:12 +02:00
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
This commit is contained in:
parent
c479ca0f44
commit
bf354b776c
2 changed files with 276 additions and 37 deletions
|
|
@ -3010,6 +3010,7 @@ class Table(Queryable):
|
||||||
num_records_processed,
|
num_records_processed,
|
||||||
replace,
|
replace,
|
||||||
ignore,
|
ignore,
|
||||||
|
list_mode=False,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Given a list ``chunk`` of records that should be written to *this* table,
|
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
|
# Build a row-list ready for executemany-style flattening
|
||||||
values = []
|
values = []
|
||||||
|
|
||||||
for record in chunk:
|
if list_mode:
|
||||||
record_values = []
|
# In list mode, records are already lists of values
|
||||||
for key in all_columns:
|
for record in chunk:
|
||||||
value = jsonify_if_needed(
|
record_values = []
|
||||||
record.get(
|
for i, key in enumerate(all_columns):
|
||||||
key,
|
if i < len(record):
|
||||||
(
|
value = jsonify_if_needed(record[i])
|
||||||
None
|
else:
|
||||||
if key != hash_id
|
value = None
|
||||||
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)
|
||||||
|
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:
|
||||||
if key in extracts:
|
extract_table = extracts[key]
|
||||||
extract_table = extracts[key]
|
value = self.db[extract_table].lookup({"value": value})
|
||||||
value = self.db[extract_table].lookup({"value": value})
|
record_values.append(value)
|
||||||
record_values.append(value)
|
values.append(record_values)
|
||||||
values.append(record_values)
|
|
||||||
|
|
||||||
columns_sql = ", ".join(f"[{c}]" for c in all_columns)
|
columns_sql = ", ".join(f"[{c}]" for c in all_columns)
|
||||||
placeholder_expr = ", ".join(conversions.get(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,
|
num_records_processed,
|
||||||
replace,
|
replace,
|
||||||
ignore,
|
ignore,
|
||||||
|
list_mode=False,
|
||||||
) -> Optional[sqlite3.Cursor]:
|
) -> Optional[sqlite3.Cursor]:
|
||||||
queries_and_params = self.build_insert_queries_and_params(
|
queries_and_params = self.build_insert_queries_and_params(
|
||||||
extracts,
|
extracts,
|
||||||
|
|
@ -3171,6 +3189,7 @@ class Table(Queryable):
|
||||||
num_records_processed,
|
num_records_processed,
|
||||||
replace,
|
replace,
|
||||||
ignore,
|
ignore,
|
||||||
|
list_mode,
|
||||||
)
|
)
|
||||||
result = None
|
result = None
|
||||||
with self.db.conn:
|
with self.db.conn:
|
||||||
|
|
@ -3200,6 +3219,7 @@ class Table(Queryable):
|
||||||
num_records_processed,
|
num_records_processed,
|
||||||
replace,
|
replace,
|
||||||
ignore,
|
ignore,
|
||||||
|
list_mode,
|
||||||
)
|
)
|
||||||
|
|
||||||
result = self.insert_chunk(
|
result = self.insert_chunk(
|
||||||
|
|
@ -3216,6 +3236,7 @@ class Table(Queryable):
|
||||||
num_records_processed,
|
num_records_processed,
|
||||||
replace,
|
replace,
|
||||||
ignore,
|
ignore,
|
||||||
|
list_mode,
|
||||||
)
|
)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
|
@ -3353,17 +3374,47 @@ class Table(Queryable):
|
||||||
all_columns = []
|
all_columns = []
|
||||||
first = True
|
first = True
|
||||||
num_records_processed = 0
|
num_records_processed = 0
|
||||||
# Fix up any records with square braces in the column names
|
|
||||||
records = fix_square_braces(records)
|
# Detect if we're using list-based iteration or dict-based iteration
|
||||||
# We can only handle a max of 999 variables in a SQL insert, so
|
list_mode = False
|
||||||
# we need to adjust the batch_size down if we have too many cols
|
column_names = None
|
||||||
records = iter(records)
|
|
||||||
# Peek at first record to count its columns:
|
# 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:
|
try:
|
||||||
first_record = next(records)
|
first_record = next(records_iter)
|
||||||
except StopIteration:
|
except StopIteration:
|
||||||
return self # It was an empty list
|
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 (
|
assert (
|
||||||
num_columns <= SQLITE_MAX_VARS
|
num_columns <= SQLITE_MAX_VARS
|
||||||
), "Rows can have a maximum of {} columns".format(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():
|
if truncate and self.exists():
|
||||||
self.db.execute("DELETE FROM [{}];".format(self.name))
|
self.db.execute("DELETE FROM [{}];".format(self.name))
|
||||||
result = None
|
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)
|
chunk = list(chunk)
|
||||||
num_records_processed += len(chunk)
|
num_records_processed += len(chunk)
|
||||||
if first:
|
if first:
|
||||||
if not self.exists():
|
if not self.exists():
|
||||||
# Use the first batch to derive the table names
|
# 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:
|
if extracts:
|
||||||
for col in extracts:
|
for col in extracts:
|
||||||
if col in column_types:
|
if col in column_types:
|
||||||
|
|
@ -3399,17 +3455,24 @@ class Table(Queryable):
|
||||||
extracts=extracts,
|
extracts=extracts,
|
||||||
strict=strict,
|
strict=strict,
|
||||||
)
|
)
|
||||||
all_columns_set = set()
|
if list_mode:
|
||||||
for record in chunk:
|
# In list mode, columns are already known
|
||||||
all_columns_set.update(record.keys())
|
all_columns = list(column_names)
|
||||||
all_columns = list(sorted(all_columns_set))
|
if hash_id:
|
||||||
if hash_id:
|
all_columns.insert(0, 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:
|
else:
|
||||||
for record in chunk:
|
if not list_mode:
|
||||||
all_columns += [
|
for record in chunk:
|
||||||
column for column in record if column not in all_columns
|
all_columns += [
|
||||||
]
|
column for column in record if column not in all_columns
|
||||||
|
]
|
||||||
|
|
||||||
first = False
|
first = False
|
||||||
|
|
||||||
|
|
@ -3427,6 +3490,7 @@ class Table(Queryable):
|
||||||
num_records_processed,
|
num_records_processed,
|
||||||
replace,
|
replace,
|
||||||
ignore,
|
ignore,
|
||||||
|
list_mode,
|
||||||
)
|
)
|
||||||
|
|
||||||
# If we only handled a single row populate self.last_pk
|
# If we only handled a single row populate self.last_pk
|
||||||
|
|
|
||||||
175
tests/test_list_mode.py
Normal file
175
tests/test_list_mode.py
Normal file
|
|
@ -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}
|
||||||
Loading…
Add table
Add a link
Reference in a new issue