From 9a6e7dab4d1f95d82a3468959d14575c0f265a56 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 21 Sep 2020 19:15:43 -0700
Subject: [PATCH] .transform() can now change the primary key

---
 sqlite_utils/db.py      | 13 +++++++------
 tests/test_transform.py | 20 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index df56080..cfe8c01 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -734,7 +734,7 @@ class Table(Queryable):
             columns=columns,
             rename=rename,
             drop=None,
-            # pk=pk,
+            pk=pk,
             # foreign_keys=foreign_keys,
             # column_order=column_order,
             # not_null=not_null,
@@ -752,6 +752,7 @@ class Table(Queryable):
         columns=None,
         rename=None,
         drop=None,
+        pk=None,
         tmp_suffix=None,
     ):
         columns = columns or {}
@@ -773,11 +774,11 @@ class Table(Queryable):
             copy_from_to[name] = new_name
 
         sqls = []
-        pk = None
-        if len(self.pks) == 1:
-            pk = self.pks[0]
-        else:
-            pk = self.pks
+        if pk is None:
+            if len(self.pks) == 1:
+                pk = self.pks[0]
+            else:
+                pk = self.pks
         sqls.append(
             self.db.create_table_sql(
                 new_table_name,
diff --git a/tests/test_transform.py b/tests/test_transform.py
index f89ca0e..06418e0 100644
--- a/tests/test_transform.py
+++ b/tests/test_transform.py
@@ -54,6 +54,26 @@ import pytest
                 "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs]",
             ],
         ),
+        # Change primary key
+        (
+            {"pk": "age"},
+            [
+                "CREATE TABLE [dogs_new_suffix] (\n   [id] INTEGER,\n   [name] TEXT,\n   [age] TEXT PRIMARY KEY\n);",
+                "INSERT INTO [dogs_new_suffix] ([id], [name], [age]) SELECT [id], [name], [age] FROM [dogs]",
+                "DROP TABLE [dogs]",
+                "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs]",
+            ],
+        ),
+        # Change primary key to a compound pk
+        (
+            {"pk": ("age", "name")},
+            [
+                "CREATE TABLE [dogs_new_suffix] (\n   [id] INTEGER,\n   [name] TEXT,\n   [age] TEXT,\n   PRIMARY KEY ([age], [name])\n);",
+                "INSERT INTO [dogs_new_suffix] ([id], [name], [age]) SELECT [id], [name], [age] FROM [dogs]",
+                "DROP TABLE [dogs]",
+                "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs]",
+            ],
+        ),
     ],
 )
 def test_transform_sql(fresh_db, params, expected_sql):