mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-26 19:04:32 +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:
Simon Willison
parent
c479ca0f44
commit
bf354b776c
2 changed files with 276 additions and 37 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue