Initial tests for .create(..., transform=True), refs #467

This commit is contained in:
Simon Willison 2022-08-27 15:11:59 -07:00
commit 1e7959d7a6
2 changed files with 79 additions and 8 deletions

View file

@ -897,7 +897,8 @@ class Database:
# Transform table to match the new definition if table already exists:
if transform and self[name].exists():
table = cast(Table, self[name])
# First add missing columns and columns to drop
needs_transform = False
# First add missing columns and figure out columns to drop
existing_columns = table.columns_dict
missing_columns = dict(
(col_name, col_type)
@ -910,13 +911,26 @@ class Database:
if missing_columns:
for col_name, col_type in missing_columns.items():
table.add_column(col_name, col_type)
# Do we need to reset the column order?
column_order = None
if list(existing_columns) != list(columns):
column_order = list(columns)
if missing_columns or columns_to_drop:
needs_transform = True
# Do we need to change the column order?
if (
column_order
and list(existing_columns)[: len(column_order)] != column_order
):
needs_transform = True
# Has the primary key changed?
current_pks = table.pks
desired_pk = None
if isinstance(pk, str):
desired_pk = [pk]
elif pk:
desired_pk = list(pk)
if desired_pk and current_pks != desired_pk:
needs_transform = True
# Only run .transform() if there is something to do
# TODO: this misses changes like pk= without also column changes
if columns_to_drop or missing_columns or column_order:
# TODO: what about not null and defaults?
if needs_transform:
table.transform(
types=columns,
drop=columns_to_drop,
@ -1692,7 +1706,9 @@ class Table(Queryable):
elif not not_null:
pass
else:
assert False, "not_null must be a dict or a set or None, it was {}".format(repr(not_null))
assert False, "not_null must be a dict or a set or None, it was {}".format(
repr(not_null)
)
# defaults=
create_table_defaults = {
(rename.get(c.name) or c.name): c.default_value

View file

@ -1155,3 +1155,58 @@ def test_create_if_no_columns(fresh_db):
with pytest.raises(AssertionError) as error:
fresh_db["t"].create({})
assert error.value.args[0] == "Tables must have at least one column"
@pytest.mark.parametrize(
"cols,kwargs,expected_schema,should_transform",
(
# Change nothing
(
{"id": int, "name": str},
{"pk": "id"},
"CREATE TABLE [demo] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT\n)",
False,
),
# Drop name column, remove primary key
({"id": int}, {}, 'CREATE TABLE "demo" (\n [id] INTEGER\n)', True),
# Add a new column
(
{"id": int, "name": str, "age": int},
{"pk": "id"},
'CREATE TABLE "demo" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] INTEGER\n)',
True,
),
# Change the primary key
(
{"id": int, "name": str},
{"pk": "name"},
'CREATE TABLE "demo" (\n [id] INTEGER,\n [name] TEXT PRIMARY KEY\n)',
True,
),
# Change in column order
(
{"id": int, "name": str},
{"pk": "id", "column_order": ["name"]},
'CREATE TABLE "demo" (\n [name] TEXT,\n [id] INTEGER PRIMARY KEY\n)',
True,
),
# Same column order is ignored
(
{"id": int, "name": str},
{"pk": "id", "column_order": ["id", "name"]},
"CREATE TABLE [demo] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT\n)",
False,
),
),
)
def test_create_transform(fresh_db, cols, kwargs, expected_schema, should_transform):
fresh_db.create_table("demo", {"id": int, "name": str}, pk="id")
fresh_db["demo"].insert({"id": 1, "name": "Cleo"})
traces = []
with fresh_db.tracer(lambda sql, parameters: traces.append((sql, parameters))):
fresh_db["demo"].create(cols, **kwargs, transform=True)
new_schema = fresh_db["demo"].schema
assert new_schema == expected_schema, repr(new_schema)
at_least_one_create_table = any(sql.startswith("CREATE TABLE") for sql, _ in traces)
assert should_transform == at_least_one_create_table
assert fresh_db["demo"].count == 1