mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-26 19:04:32 +02:00
parent
632f4f5f0e
commit
975b6c6191
2 changed files with 106 additions and 8 deletions
|
|
@ -3033,14 +3033,18 @@ class Table(Queryable):
|
|||
# Pad short records with None, truncate long ones
|
||||
record_len = len(record)
|
||||
if record_len < num_columns:
|
||||
record_values = [jsonify_if_needed(v) for v in record] + [None] * (num_columns - record_len)
|
||||
record_values = [jsonify_if_needed(v) for v in record] + [None] * (
|
||||
num_columns - record_len
|
||||
)
|
||||
else:
|
||||
record_values = [jsonify_if_needed(v) for v in record[:num_columns]]
|
||||
# Only process extracts if there are any
|
||||
if has_extracts:
|
||||
for i, key in enumerate(all_columns):
|
||||
if key in extracts:
|
||||
record_values[i] = self.db[extracts[key]].lookup({"value": record_values[i]})
|
||||
record_values[i] = self.db[extracts[key]].lookup(
|
||||
{"value": record_values[i]}
|
||||
)
|
||||
values.append(record_values)
|
||||
else:
|
||||
# Dict mode: original logic
|
||||
|
|
@ -3393,12 +3397,14 @@ class Table(Queryable):
|
|||
return self # It was an empty list
|
||||
|
||||
# Check if this is list mode or dict mode
|
||||
if isinstance(first_record, list):
|
||||
# List mode: first record should be column names
|
||||
if isinstance(first_record, (list, tuple)):
|
||||
# List/tuple 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
|
||||
raise ValueError(
|
||||
"When using list-based iteration, the first yielded value must be a list of column name strings"
|
||||
)
|
||||
column_names = list(first_record)
|
||||
all_columns = column_names
|
||||
num_columns = len(column_names)
|
||||
# Get the actual first data record
|
||||
|
|
@ -3406,8 +3412,10 @@ class Table(Queryable):
|
|||
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")
|
||||
if not isinstance(first_record, (list, tuple)):
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""
|
||||
Tests for list-based iteration in insert_all and upsert_all
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from sqlite_utils import Database
|
||||
|
||||
|
|
@ -173,3 +174,92 @@ def test_backwards_compatibility_dict_mode():
|
|||
rows = list(db["people"].rows)
|
||||
assert len(rows) == 2
|
||||
assert rows[0] == {"id": 1, "name": "Alice", "age": 30}
|
||||
|
||||
|
||||
def test_insert_all_tuple_mode_basic():
|
||||
"""Test basic insert_all with tuple-based iteration"""
|
||||
db = Database(memory=True)
|
||||
|
||||
def data_generator():
|
||||
# First yield column names as tuple
|
||||
yield ("id", "name", "age")
|
||||
# Then yield data rows as tuples
|
||||
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_mixed_list_tuple():
|
||||
"""Test insert_all with mixed lists and tuples for data rows"""
|
||||
db = Database(memory=True)
|
||||
|
||||
def data_generator():
|
||||
# Column names as list
|
||||
yield ["id", "name", "age"]
|
||||
# Mix of list and tuple data rows
|
||||
yield [1, "Alice", 30]
|
||||
yield (2, "Bob", 25)
|
||||
yield [3, "Charlie", 35]
|
||||
yield (4, "Diana", 40)
|
||||
|
||||
db["people"].insert_all(data_generator())
|
||||
|
||||
rows = list(db["people"].rows)
|
||||
assert len(rows) == 4
|
||||
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}
|
||||
assert rows[3] == {"id": 4, "name": "Diana", "age": 40}
|
||||
|
||||
|
||||
def test_upsert_all_tuple_mode():
|
||||
"""Test upsert_all with tuple-based iteration"""
|
||||
db = Database(memory=True)
|
||||
|
||||
# Initial insert with tuples
|
||||
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 tuples
|
||||
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_tuple_mode_shorter_rows():
|
||||
"""Test that tuple 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}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue