From 987dd123f2ac43c5ab66d69e59d454fe09660606 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Mon, 21 Sep 2020 21:20:01 -0700 Subject: [PATCH] table.transform() method - closes #114 --- docs/python-api.rst | 93 +++++++++++++ sqlite_utils/db.py | 179 +++++++++++++++++++++++- tests/test_transform.py | 298 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 568 insertions(+), 2 deletions(-) create mode 100644 tests/test_transform.py diff --git a/docs/python-api.rst b/docs/python-api.rst index 2ded7b0..026e2e0 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -908,6 +908,99 @@ You can drop a table or view using the ``.drop()`` method: db["my_table"].drop() +.. _python_api_transform: + +Transforming a table +==================== + +The SQLite ``ALTER TABLE`` statement is limited. It can add columns and rename tables, but it cannot rename columns, drop columns, change column types, change ``NOT NULL`` status or change the primary key for a table. + +The ``table.transform()`` method can do all of these things, by implementing a multi-step pattern `described in the SQLite documentation `__: + +1. Start a transaction +2. ``CREATE TABLE tablename_new_x123`` with the required changes +3. Copy the old data into the new table using ``INSERT INTO tablename_new_x123 SELECT * FROM tablename;`` +4. ``DROP TABLE tablename;`` +5. ``ALTER TABLE tablename_new_x123 RENAME TO tablename;`` +6. Commit the transaction + +The ``.transform()`` method takes a number of parameters, all of which are optional. + +To alter the type of a column, use the first argument: + +.. code-block:: python + + # Convert the 'age' column to an integer, and 'weight' to a float + table.transform({"age": int, "weight": float}) + +The ``rename=`` parameter can rename columns: + +.. code-block:: python + + # Rename 'age' to 'initial_age': + table.transform(rename={"age": "initial_age"}) + +To drop columns, pass them in the ``drop=`` set: + +.. code-block:: python + + # Drop the 'age' column: + table.transform(drop={"age"}) + +To change the primary key for a table, use ``pk=``. This can be passed a single column for a regular primary key, or a tuple of columns to create a compound primary key. Passing ``pk=None`` will remove the primary key and convert the table into a ``rowid`` table. + +.. code-block:: python + + # Make `user_id` the new primary key + table.transform(pk="user_id") + +You can change the ``NOT NULL`` status of columns by using ``not_null=``. You can pass this a set of columns to make those columns ``NOT NULL``: + +.. code-block:: python + + # Make the 'age' and 'weight' columns NOT NULL + table.transform(not_null={"age", "weight"}) + +If you want to take existing ``NOT NULL`` columns and change them to allow null values, you can do so by passing a dictionary of true/false values instead: + +.. code-block:: python + + # 'age' is NOT NULL but we want to allow NULL: + table.transform(not_null={"age": False}) + + # Make age allow NULL and switch weight to being NOT NULL: + table.transform(not_null={"age": False, "weight": True}) + +The ``defaults=`` parameter can be used to set or change the defaults for different columns: + +.. code-block:: python + + # Set default age to 1: + table.transform(defaults={"age": 1}) + + # Now remove the default from that column: + table.transform(defaults={"age": None}) + +You can use ``.transform()`` to remove foreign key constraints from a table. You will need to know the name of the column, the name of the table it points to and the name of the column it references on that other table. + +This example drops two foreign keys - the one from ``places.country`` to ``country.id`` and the one from ``places.continent`` to ``continent.id``: + +.. code-block:: python + + db["places"].transform( + drop_foreign_keys=( + ("country", "country", "id"), + ("continent", "continent", "id"), + ) + ) + +.. _python_api_transform_sql: + +Custom transformations with .transform_sql() +-------------------------------------------- + +If you want to do something more advanced, you can call the ``table.transform_sql(...)`` method with the same arguments that you would have passed to ``table.transform(...)``. This method will return a list of SQL statements that should be executed to implement the change. You can then make modifications to that SQL before executing it yourself. + .. _python_api_hash: Setting an ID based on the hash of the row contents diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index daf3e25..d8ad05b 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -268,7 +268,7 @@ class Database: ) return fks - def create_table( + def create_table_sql( self, name, columns, @@ -336,7 +336,7 @@ class Database: column_extras.append("PRIMARY KEY") if column_name in not_null: column_extras.append("NOT NULL") - if column_name in defaults: + if column_name in defaults and defaults[column_name] is not None: column_extras.append( "DEFAULT {}".format(self.escape(defaults[column_name])) ) @@ -368,6 +368,31 @@ class Database: """.format( table=name, columns_sql=columns_sql, extra_pk=extra_pk ) + return sql + + def create_table( + self, + name, + columns, + pk=None, + foreign_keys=None, + column_order=None, + not_null=None, + defaults=None, + hash_id=None, + extracts=None, + ): + sql = self.create_table_sql( + name=name, + columns=columns, + pk=pk, + foreign_keys=foreign_keys, + column_order=column_order, + not_null=not_null, + defaults=defaults, + hash_id=hash_id, + extracts=extracts, + ) self.execute(sql) return self.table( name, @@ -691,6 +716,156 @@ class Table(Queryable): ) return self + def transform( + self, + columns=None, + rename=None, + drop=None, + pk=DEFAULT, + not_null=None, + defaults=None, + drop_foreign_keys=None, + ): + assert self.exists(), "Cannot transform a table that doesn't exist yet" + sqls = self.transform_sql( + columns=columns, + rename=rename, + drop=None, + pk=pk, + not_null=not_null, + defaults=defaults, + drop_foreign_keys=drop_foreign_keys, + ) + initial_pragma_foreign_keys = self.db.execute("PRAGMA foreign_keys").fetchone()[ + 0 + ] + try: + with self.db.conn: + for sql in sqls: + self.db.execute(sql) + finally: + # Make sure we reset PRAGMA foreign_keys correctly + if ( + initial_pragma_foreign_keys + and not self.db.execute("PRAGMA foreign_keys").fetchone()[0] + ): + self.db.execute("PRAGMA foreign_keys=1") + return self + + def transform_sql( + self, + columns=None, + rename=None, + drop=None, + pk=DEFAULT, + not_null=None, + defaults=None, + drop_foreign_keys=None, + tmp_suffix=None, + ): + columns = columns or {} + rename = rename or {} + drop = drop or set() + new_table_name = "{}_new_{}".format( + self.name, tmp_suffix or os.urandom(6).hex() + ) + current_column_pairs = list(self.columns_dict.items()) + new_column_pairs = [] + copy_from_to = {column: column for column, _ in current_column_pairs} + for name, type_ in current_column_pairs: + type_ = columns.get(name) or type_ + if name in drop: + del [copy_from_to[name]] + continue + new_name = rename.get(name) or name + new_column_pairs.append((new_name, type_)) + copy_from_to[name] = new_name + + should_flip_foreign_keys_pragma = self.db.execute( + "PRAGMA foreign_keys" + ).fetchone()[0] + + sqls = [] + + if should_flip_foreign_keys_pragma: + sqls.append("PRAGMA foreign_keys=OFF") + + if pk is DEFAULT: + pks_renamed = tuple(rename.get(p) or p for p in self.pks) + if len(pks_renamed) == 1: + pk = pks_renamed[0] + else: + pk = pks_renamed + + # not_null may be a set or dict, need to convert to a set + create_table_not_null = {c.name for c in self.columns if c.notnull} + if isinstance(not_null, dict): + # Remove any columns with a value of False + for key, value in not_null.items(): + # Column may have been renamed + key = rename.get(key) or key + if value is False and key in create_table_not_null: + create_table_not_null.remove(key) + else: + create_table_not_null.add(key) + elif isinstance(not_null, set): + create_table_not_null.update(rename.get(k) or k for k in not_null) + + # defaults= + create_table_defaults = { + (rename.get(c.name) or c.name): c.default_value + for c in self.columns + if c.default_value is not None + } + if defaults is not None: + create_table_defaults.update( + {rename.get(c) or c: v for c, v in defaults.items()} + ) + + # foreign_keys + create_table_foreign_keys = [] + for table, column, other_table, other_column in self.foreign_keys: + if (drop_foreign_keys is None) or ( + (column, other_table, other_column) not in drop_foreign_keys + ): + create_table_foreign_keys.append( + (rename.get(column) or column, other_table, other_column) + ) + + sqls.append( + self.db.create_table_sql( + new_table_name, + dict(new_column_pairs), + pk=pk, + not_null=create_table_not_null, + defaults=create_table_defaults, + foreign_keys=create_table_foreign_keys, + ).strip() + ) + # Copy across data, respecting any renamed columns + new_cols = [] + old_cols = [] + for from_, to_ in copy_from_to.items(): + old_cols.append(from_) + new_cols.append(to_) + copy_sql = "INSERT INTO [{new_table}] ({new_cols}) SELECT {old_cols} FROM [{old_table}]".format( + new_table=new_table_name, + old_table=self.name, + old_cols=", ".join("[{}]".format(col) for col in old_cols), + 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)) + # Rename the new one + sqls.append("ALTER TABLE [{}] RENAME TO [{}]".format(new_table_name, self.name)) + + if should_flip_foreign_keys_pragma: + sqls.append("PRAGMA foreign_key_check") + sqls.append("PRAGMA foreign_keys=ON") + + return sqls + def create_index(self, columns, index_name=None, unique=False, if_not_exists=False): if index_name is None: index_name = "idx_{}_{}".format( diff --git a/tests/test_transform.py b/tests/test_transform.py new file mode 100644 index 0000000..c49efe1 --- /dev/null +++ b/tests/test_transform.py @@ -0,0 +1,298 @@ +from sqlite_utils.db import ForeignKey +from sqlite_utils.utils import OperationalError +import pytest + + +@pytest.mark.parametrize( + "params,expected_sql", + [ + # Identity transform - nothing changes + ( + {}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] TEXT\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 column type + ( + {"columns": {"age": int}}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] INTEGER\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]", + ], + ), + # Rename a column + ( + {"rename": {"age": "dog_age"}}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [dog_age] TEXT\n);", + "INSERT INTO [dogs_new_suffix] ([id], [name], [dog_age]) SELECT [id], [name], [age] FROM [dogs]", + "DROP TABLE [dogs]", + "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs]", + ], + ), + # Drop a column + ( + {"drop": ["age"]}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT\n);", + "INSERT INTO [dogs_new_suffix] ([id], [name]) SELECT [id], [name] FROM [dogs]", + "DROP TABLE [dogs]", + "ALTER TABLE [dogs_new_suffix] RENAME TO [dogs]", + ], + ), + # Convert type AND rename column + ( + {"columns": {"age": int}, "rename": {"age": "dog_age"}}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [dog_age] INTEGER\n);", + "INSERT INTO [dogs_new_suffix] ([id], [name], [dog_age]) SELECT [id], [name], [age] FROM [dogs]", + "DROP TABLE [dogs]", + "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]", + ], + ), + # Remove primary key, creating a rowid table + ( + {"pk": None}, + [ + "CREATE TABLE [dogs_new_suffix] (\n [id] INTEGER,\n [name] TEXT,\n [age] TEXT\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]", + ], + ), + ], +) +@pytest.mark.parametrize("use_pragma_foreign_keys", [False, True]) +def test_transform_sql(fresh_db, params, expected_sql, use_pragma_foreign_keys): + dogs = fresh_db["dogs"] + if use_pragma_foreign_keys: + fresh_db.conn.execute("PRAGMA foreign_keys=ON") + expected_sql.insert(0, "PRAGMA foreign_keys=OFF") + expected_sql.append("PRAGMA foreign_key_check") + expected_sql.append("PRAGMA foreign_keys=ON") + dogs.insert({"id": 1, "name": "Cleo", "age": "5"}, pk="id") + sql = dogs.transform_sql(**{**params, **{"tmp_suffix": "suffix"}}) + assert sql == expected_sql + # Check that .transform() runs without exceptions: + dogs.transform(**params) + + +def test_transform_sql_rowid_to_id(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo", "age": "5"}) + assert ( + dogs.schema + == "CREATE TABLE [dogs] (\n [id] INTEGER,\n [name] TEXT,\n [age] TEXT\n)" + ) + dogs.transform(pk="id") + # Slight oddity: [dogs] becomes "dogs" during the rename: + assert ( + dogs.schema + == 'CREATE TABLE "dogs" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] TEXT\n)' + ) + + +def test_transform_rename_pk(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo", "age": "5"}, pk="id") + dogs.transform(rename={"id": "pk"}) + assert ( + dogs.schema + == 'CREATE TABLE "dogs" (\n [pk] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] TEXT\n)' + ) + + +def test_transform_not_null(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo", "age": "5"}, pk="id") + dogs.transform(not_null={"name"}) + assert ( + dogs.schema + == 'CREATE TABLE "dogs" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT NOT NULL,\n [age] TEXT\n)' + ) + + +def test_transform_remove_a_not_null(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo", "age": "5"}, not_null={"age"}, pk="id") + dogs.transform(not_null={"name": True, "age": False}) + assert ( + dogs.schema + == 'CREATE TABLE "dogs" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT NOT NULL,\n [age] TEXT\n)' + ) + + +@pytest.mark.parametrize("not_null", [{"age"}, {"age": True}]) +def test_transform_add_not_null_with_rename(fresh_db, not_null): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo", "age": "5"}, pk="id") + dogs.transform(not_null=not_null, rename={"age": "dog_age"}) + assert ( + dogs.schema + == 'CREATE TABLE "dogs" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [dog_age] TEXT NOT NULL\n)' + ) + + +def test_transform_defaults(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo", "age": 5}, pk="id") + dogs.transform(defaults={"age": 1}) + assert ( + dogs.schema + == 'CREATE TABLE "dogs" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] INTEGER DEFAULT 1\n)' + ) + + +def test_transform_defaults_and_rename_column(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo", "age": 5}, pk="id") + dogs.transform(rename={"age": "dog_age"}, defaults={"age": 1}) + assert ( + dogs.schema + == 'CREATE TABLE "dogs" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [dog_age] INTEGER DEFAULT 1\n)' + ) + + +def test_remove_defaults(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo", "age": 5}, defaults={"age": 1}, pk="id") + dogs.transform(defaults={"age": None}) + assert ( + dogs.schema + == 'CREATE TABLE "dogs" (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT,\n [age] INTEGER\n)' + ) + + +@pytest.fixture +def authors_db(fresh_db): + books = fresh_db["books"] + authors = fresh_db["authors"] + authors.insert({"id": 5, "name": "Jane McGonical"}, pk="id") + books.insert( + {"id": 2, "title": "Reality is Broken", "author_id": 5}, + foreign_keys=("author_id",), + pk="id", + ) + return fresh_db + + +def test_transform_foreign_keys_persist(authors_db): + assert authors_db["books"].foreign_keys == [ + ForeignKey( + table="books", column="author_id", other_table="authors", other_column="id" + ) + ] + authors_db["books"].transform(rename={"title": "book_title"}) + assert authors_db["books"].foreign_keys == [ + ForeignKey( + table="books", column="author_id", other_table="authors", other_column="id" + ) + ] + + +@pytest.mark.parametrize("use_pragma_foreign_keys", [False, True]) +def test_transform_foreign_keys_survive_renamed_column( + authors_db, use_pragma_foreign_keys +): + if use_pragma_foreign_keys: + authors_db.conn.execute("PRAGMA foreign_keys=ON") + authors_db["books"].transform(rename={"author_id": "author_id_2"}) + assert authors_db["books"].foreign_keys == [ + ForeignKey( + table="books", + column="author_id_2", + other_table="authors", + other_column="id", + ) + ] + + +@pytest.mark.parametrize("use_pragma_foreign_keys", [False, True]) +def test_transform_drop_foreign_keys(fresh_db, use_pragma_foreign_keys): + if use_pragma_foreign_keys: + fresh_db.conn.execute("PRAGMA foreign_keys=ON") + # Create table with three foreign keys so we can drop two of them + fresh_db["country"].insert({"id": 1, "name": "France"}, pk="id") + fresh_db["continent"].insert({"id": 2, "name": "Europe"}, pk="id") + fresh_db["city"].insert({"id": 24, "name": "Paris"}, pk="id") + fresh_db["places"].insert( + { + "id": 32, + "name": "Caveau de la Huchette", + "country": 1, + "continent": 2, + "city": 24, + }, + foreign_keys=("country", "continent", "city"), + ) + assert fresh_db["places"].foreign_keys == [ + ForeignKey( + table="places", column="city", other_table="city", other_column="id" + ), + ForeignKey( + table="places", + column="continent", + other_table="continent", + other_column="id", + ), + ForeignKey( + table="places", column="country", other_table="country", other_column="id" + ), + ] + # Drop two of those foreign keys + fresh_db["places"].transform( + drop_foreign_keys=( + ("country", "country", "id"), + ("continent", "continent", "id"), + ) + ) + # Should be only one foreign key now + assert fresh_db["places"].foreign_keys == [ + ForeignKey(table="places", column="city", other_table="city", other_column="id") + ] + if use_pragma_foreign_keys: + assert fresh_db.conn.execute("PRAGMA foreign_keys").fetchone()[0] + + +def test_transform_verify_foreign_keys(fresh_db): + fresh_db.conn.execute("PRAGMA foreign_keys=ON") + fresh_db["authors"].insert({"id": 3, "name": "Tina"}, pk="id") + fresh_db["books"].insert( + {"id": 1, "title": "Book", "author_id": 3}, pk="id", foreign_keys={"author_id"} + ) + # Renaming the id column on authors should break everything + with pytest.raises(OperationalError) as e: + fresh_db["authors"].transform(rename={"id": "id2"}) + assert e.value.args[0] == 'foreign key mismatch - "books" referencing "authors"' + # This should have rolled us back + assert ( + fresh_db["authors"].schema + == "CREATE TABLE [authors] (\n [id] INTEGER PRIMARY KEY,\n [name] TEXT\n)" + ) + assert fresh_db.conn.execute("PRAGMA foreign_keys").fetchone()[0]