From 1e7959d7a68ae5faabaabfcfac505c82fa426901 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 27 Aug 2022 15:11:59 -0700
Subject: [PATCH] Initial tests for .create(..., transform=True), refs #467

---
 sqlite_utils/db.py   | 32 +++++++++++++++++++-------
 tests/test_create.py | 55 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 25fded0..49b1f68 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -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
diff --git a/tests/test_create.py b/tests/test_create.py
index 60181de..6909e6f 100644
--- a/tests/test_create.py
+++ b/tests/test_create.py
@@ -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