From 61aaa698154c9020934d4896d04174a0f9c4f777 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 22 Jul 2023 15:32:09 -0700
Subject: [PATCH] .transform(..., keep_table=name) parameter, closes #571

Also type hints for the transform_sql() method
---
 docs/python-api.rst     |  6 ++++++
 sqlite_utils/db.py      | 39 ++++++++++++++++++++++++++-------------
 tests/test_transform.py | 10 ++++++++++
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/docs/python-api.rst b/docs/python-api.rst
index ba787d2..21fad3d 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -1376,6 +1376,12 @@ The ``.transform()`` method takes a number of parameters, all of which are optio
 
 As a bonus, calling ``.transform()`` will reformat the schema for the table that is stored in SQLite to make it more readable. This works even if you call it without any arguments.
 
+To keep the original table around instead of dropping it, pass the ``keep_table=`` option and specify the name of the table you would like it to be renamed to:
+
+.. code-block:: python
+
+    table.transform(types={"age": int}, keep_table="original_table")
+
 Altering column types
 ---------------------
 
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 7098ec0..663a184 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -1702,6 +1702,7 @@ class Table(Queryable):
         defaults: Optional[Dict[str, Any]] = None,
         drop_foreign_keys: Optional[Iterable] = None,
         column_order: Optional[List[str]] = None,
+        keep_table: Optional[str] = None,
     ) -> "Table":
         """
         Apply an advanced alter table, including operations that are not supported by
@@ -1717,7 +1718,9 @@ class Table(Queryable):
         :param defaults: Default values for columns
         :param drop_foreign_keys: Names of columns that should have their foreign key constraints removed
         :param column_order: List of strings specifying a full or partial column order
-          to use when creating the table.
+          to use when creating the table
+        :param keep_table: If specified, the existing table will be renamed to this and will not be
+          dropped
         """
         assert self.exists(), "Cannot transform a table that doesn't exist yet"
         sqls = self.transform_sql(
@@ -1729,6 +1732,7 @@ class Table(Queryable):
             defaults=defaults,
             drop_foreign_keys=drop_foreign_keys,
             column_order=column_order,
+            keep_table=keep_table,
         )
         pragma_foreign_keys_was_on = self.db.execute("PRAGMA foreign_keys").fetchone()[
             0
@@ -1750,15 +1754,16 @@ class Table(Queryable):
     def transform_sql(
         self,
         *,
-        types=None,
-        rename=None,
-        drop=None,
-        pk=DEFAULT,
-        not_null=None,
-        defaults=None,
-        drop_foreign_keys=None,
-        column_order=None,
-        tmp_suffix=None,
+        types: Optional[dict] = None,
+        rename: Optional[dict] = None,
+        drop: Optional[Iterable] = None,
+        pk: Optional[Any] = DEFAULT,
+        not_null: Optional[Iterable[str]] = None,
+        defaults: Optional[Dict[str, Any]] = None,
+        drop_foreign_keys: Optional[Iterable] = None,
+        column_order: Optional[List[str]] = None,
+        tmp_suffix: Optional[str] = None,
+        keep_table: Optional[str] = None,
     ) -> List[str]:
         """
         Return a list of SQL statements that should be executed in order to apply this transformation.
@@ -1771,7 +1776,10 @@ class Table(Queryable):
         :param defaults: Default values for columns
         :param drop_foreign_keys: Names of columns that should have their foreign key constraints removed
         :param column_order: List of strings specifying a full or partial column order
-          to use when creating the table.
+          to use when creating the table
+        :param tmp_suffix: Suffix to use for the temporary table name
+        :param keep_table: If specified, the existing table will be renamed to this and will not be
+          dropped
         """
         types = types or {}
         rename = rename or {}
@@ -1872,8 +1880,13 @@ class Table(Queryable):
             new_cols=", ".join("[{}]".format(col) for col in new_cols),
         )
         sqls.append(copy_sql)
-        # Drop the old table
-        sqls.append("DROP TABLE [{}];".format(self.name))
+        # Drop (or keep) the old table
+        if keep_table:
+            sqls.append(
+                "ALTER TABLE [{}] RENAME TO [{}];".format(self.name, keep_table)
+            )
+        else:
+            sqls.append("DROP TABLE [{}];".format(self.name))
         # Rename the new one
         sqls.append(
             "ALTER TABLE [{}] RENAME TO [{}];".format(new_table_name, self.name)
diff --git a/tests/test_transform.py b/tests/test_transform.py
index 06e5729..6b0b53a 100644
--- a/tests/test_transform.py
+++ b/tests/test_transform.py
@@ -86,6 +86,16 @@ import pytest
                 "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs];",
             ],
         ),
+        # Keeping the table
+        (
+            {"drop": ["age"], "keep_table": "kept_table"},
+            [
+                "CREATE TABLE [dogs_new_suffix] (\n   [id] INTEGER PRIMARY KEY,\n   [name] TEXT\n);",
+                "INSERT INTO [dogs_new_suffix] ([id], [name])\n   SELECT [id], [name] FROM [dogs];",
+                "ALTER TABLE [dogs] RENAME TO [kept_table];",
+                "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs];",
+            ],
+        ),
     ],
 )
 @pytest.mark.parametrize("use_pragma_foreign_keys", [False, True])