Merge pull request piccolo-orm#71 from piccolo-orm/migration_test_coverage

dantownsend · web-flow · commit 3969a252470e · 2021-02-14T16:34:52.000Z
Migration test coverage
diff --git a/piccolo/apps/migrations/auto/diffable_table.py b/piccolo/apps/migrations/auto/diffable_table.py
@@ -102,7 +102,7 @@ def __sub__(self, value: DiffableTable) -> TableDelta:
                 for key, _ in delta.items()
             }
 
-            if delta:
+            if delta or (column.__class__ != existing_column.__class__):
                 alter_columns.append(
                     AlterColumn(
                         table_class_name=self.class_name,
diff --git a/piccolo/query/methods/alter.py b/piccolo/query/methods/alter.py
@@ -93,20 +93,30 @@ def querystring(self) -> QueryString:
 
 @dataclass
 class SetColumnType(AlterStatement):
-    __slots__ = ("old_column", "new_column")
+    """
+    :param using_expression:
+        Postgres can't automatically convert between certain column types. You
+        can tell Postgres which action to take. For example
+        `my_column_name::integer`.
+
+    """
 
     old_column: Column
     new_column: Column
+    using_expression: t.Optional[str] = None
 
     @property
     def querystring(self) -> QueryString:
         if self.new_column._meta._table is None:
             self.new_column._meta._table = self.old_column._meta.table
 
         column_name = self.old_column._meta.name
-        return QueryString(
+        query = (
             f"ALTER COLUMN {column_name} TYPE {self.new_column.column_type}"
         )
+        if self.using_expression is not None:
+            query += f" USING {self.using_expression}"
+        return QueryString(query)
 
 
 @dataclass
@@ -353,12 +363,21 @@ def rename_column(
         self._rename_columns.append(RenameColumn(column, new_name))
         return self
 
-    def set_column_type(self, old_column: Column, new_column: Column) -> Alter:
+    def set_column_type(
+        self,
+        old_column: Column,
+        new_column: Column,
+        using_expression: t.Optional[str] = None,
+    ) -> Alter:
         """
         Change the type of a column.
         """
         self._set_column_type.append(
-            SetColumnType(old_column=old_column, new_column=new_column)
+            SetColumnType(
+                old_column=old_column,
+                new_column=new_column,
+                using_expression=using_expression,
+            )
         )
         return self
 
diff --git a/tests/apps/migrations/auto/test_migration_manager.py b/tests/apps/migrations/auto/test_migration_manager.py
@@ -4,7 +4,7 @@
 from asyncpg.exceptions import UniqueViolationError
 from piccolo.apps.migrations.auto import MigrationManager
 from piccolo.apps.migrations.commands.base import BaseMigrationManager
-from piccolo.columns import Varchar
+from piccolo.columns import Varchar, Text
 from piccolo.columns.base import OnDelete, OnUpdate
 
 from tests.example_app.tables import Manager
@@ -146,6 +146,36 @@ def test_add_column(self):
         response = self.run_sync("SELECT * FROM manager;")
         self.assertEqual(response, [{"id": 1, "name": "Dave"}])
 
+    @postgres_only
+    def test_add_column_with_index(self):
+        """
+        Test adding a column with an index to a MigrationManager.
+        """
+        manager = MigrationManager()
+        manager.add_column(
+            table_class_name="Manager",
+            tablename="manager",
+            column_name="email",
+            column_class_name="Varchar",
+            params={
+                "length": 100,
+                "default": "",
+                "null": True,
+                "primary": False,
+                "key": False,
+                "unique": True,
+                "index": True,
+            },
+        )
+        index_name = Manager._get_index_name(["email"])
+
+        asyncio.run(manager.run())
+        self.assertTrue(index_name in Manager.indexes().run_sync())
+
+        # Reverse
+        asyncio.run(manager.run_backwards())
+        self.assertTrue(index_name not in Manager.indexes().run_sync())
+
     @postgres_only
     def test_add_foreign_key_self_column(self):
         """
@@ -310,6 +340,36 @@ def test_alter_column_unique(self):
         response = self.run_sync("SELECT name FROM manager;")
         self.assertEqual(response, [{"name": "Dave"}, {"name": "Dave"}])
 
+    @postgres_only
+    def test_alter_column_set_null(self):
+        """
+        Test altering whether a column is nullable with MigrationManager.
+        """
+        manager = MigrationManager()
+
+        manager.alter_column(
+            table_class_name="Manager",
+            tablename="manager",
+            column_name="name",
+            params={"null": True},
+            old_params={"null": False},
+        )
+
+        asyncio.run(manager.run())
+        self.assertTrue(
+            self.get_postgres_is_nullable(
+                tablename="manager", column_name="name"
+            )
+        )
+
+        # Reverse
+        asyncio.run(manager.run_backwards())
+        self.assertFalse(
+            self.get_postgres_is_nullable(
+                tablename="manager", column_name="name"
+            )
+        )
+
     def _get_column_precision_and_scale(
         self, tablename="ticket", column_name="price"
     ):
@@ -467,6 +527,68 @@ def test_alter_column_add_index(self):
             not in Manager.indexes().run_sync()
         )
 
+    @postgres_only
+    def test_alter_column_set_type(self):
+        """
+        Test altering a column to change it's type with MigrationManager.
+        """
+        manager = MigrationManager()
+
+        manager.alter_column(
+            table_class_name="Manager",
+            tablename="manager",
+            column_name="name",
+            params={},
+            old_params={},
+            column_class=Text,
+            old_column_class=Varchar,
+        )
+
+        asyncio.run(manager.run())
+        column_type_str = self.get_postgres_column_type(
+            tablename="manager", column_name="name"
+        )
+        self.assertEqual(column_type_str, "TEXT")
+
+        asyncio.run(manager.run_backwards())
+        column_type_str = self.get_postgres_column_type(
+            tablename="manager", column_name="name"
+        )
+        self.assertEqual(column_type_str, "CHARACTER VARYING")
+
+    @postgres_only
+    def test_alter_column_set_length(self):
+        """
+        Test altering a Varchar column's length with MigrationManager.
+        """
+        manager = MigrationManager()
+
+        manager.alter_column(
+            table_class_name="Manager",
+            tablename="manager",
+            column_name="name",
+            params={"length": 500},
+            old_params={"length": 200},
+            column_class=Text,
+            old_column_class=Varchar,
+        )
+
+        asyncio.run(manager.run())
+        self.assertEqual(
+            self.get_postgres_varchar_length(
+                tablename="manager", column_name="name"
+            ),
+            500,
+        )
+
+        asyncio.run(manager.run_backwards())
+        self.assertEqual(
+            self.get_postgres_varchar_length(
+                tablename="manager", column_name="name"
+            ),
+            200,
+        )
+
     @postgres_only
     @patch.object(BaseMigrationManager, "get_migration_managers")
     def test_drop_table(self, get_migration_managers: MagicMock):
diff --git a/tests/apps/migrations/auto/test_schema_differ.py b/tests/apps/migrations/auto/test_schema_differ.py
@@ -14,7 +14,31 @@ def test_add_table(self):
         """
         Test adding a new table.
         """
-        pass
+        name_column = Varchar()
+        name_column._meta.name = "name"
+        schema: t.List[DiffableTable] = [
+            DiffableTable(
+                class_name="Band", tablename="band", columns=[name_column]
+            )
+        ]
+        schema_snapshot: t.List[DiffableTable] = []
+        schema_differ = SchemaDiffer(
+            schema=schema, schema_snapshot=schema_snapshot, auto_input="y"
+        )
+
+        create_tables = schema_differ.create_tables
+        self.assertTrue(len(create_tables.statements) == 1)
+        self.assertEqual(
+            create_tables.statements[0],
+            "manager.add_table('Band', tablename='band')",
+        )
+
+        new_table_columns = schema_differ.new_table_columns
+        self.assertTrue(len(new_table_columns.statements) == 1)
+        self.assertEqual(
+            new_table_columns.statements[0],
+            "manager.add_column(table_class_name='Band', tablename='band', column_name='name', column_class_name='Varchar', column_class=Varchar, params={'length': 255, 'default': '', 'null': False, 'primary': False, 'key': False, 'unique': False, 'index': False})",  # noqa
+        )
 
     def test_drop_table(self):
         """
diff --git a/tests/base.py b/tests/base.py
@@ -50,11 +50,60 @@ def run_sync(self, query):
         _Table = type("_Table", (Table,), {})
         return _Table.raw(query).run_sync()
 
-    def table_exists(self, tablename: str):
+    def table_exists(self, tablename: str) -> bool:
         _Table: t.Type[Table] = type(tablename.upper(), (Table,), {})
         _Table._meta.tablename = tablename
         return _Table.table_exists().run_sync()
 
+    ###########################################################################
+
+    # Postgres specific utils
+
+    def get_postgres_column_definition(
+        self, tablename: str, column_name: str
+    ) -> t.Dict[str, t.Any]:
+        query = """
+            SELECT * FROM information_schema.columns
+            WHERE table_name = '{tablename}'
+            AND table_catalog = 'piccolo'
+            AND column_name = '{column_name}'
+        """.format(
+            tablename=tablename, column_name=column_name
+        )
+        response = self.run_sync(query)
+        return response[0]
+
+    def get_postgres_column_type(
+        self, tablename: str, column_name: str
+    ) -> str:
+        """
+        Fetches the column type as a string, from the database.
+        """
+        return self.get_postgres_column_definition(
+            tablename=tablename, column_name=column_name
+        )["data_type"].upper()
+
+    def get_postgres_is_nullable(self, tablename, column_name: str) -> bool:
+        """
+        Fetches whether the column is defined as nullable, from the database.
+        """
+        return (
+            self.get_postgres_column_definition(
+                tablename=tablename, column_name=column_name
+            )["is_nullable"].upper()
+            == "YES"
+        )
+
+    def get_postgres_varchar_length(self, tablename, column_name: str) -> int:
+        """
+        Fetches whether the column is defined as nullable, from the database.
+        """
+        return self.get_postgres_column_definition(
+            tablename=tablename, column_name=column_name
+        )["character_maximum_length"]
+
+    ###########################################################################
+
     def create_tables(self):
         if ENGINE.engine_type == "postgres":
             self.run_sync(
diff --git a/tests/table/test_alter.py b/tests/table/test_alter.py
@@ -1,7 +1,6 @@
-from piccolo.columns.column_types import Varchar
 from unittest import TestCase
 
-from piccolo.columns import BigInt, Integer, Numeric
+from piccolo.columns import BigInt, Integer, Numeric, Varchar
 from piccolo.table import Table
 
 from ..base import DBTestCase, postgres_only
@@ -151,15 +150,12 @@ def test_integer_to_bigint(self):
         )
         alter_query.run_sync()
 
-        query = """
-            SELECT data_type FROM information_schema.columns
-            WHERE table_name = 'band'
-            AND table_catalog = 'piccolo'
-            AND column_name = 'popularity'
-            """
-
-        response = Band.raw(query).run_sync()
-        self.assertEqual(response[0]["data_type"].upper(), "BIGINT")
+        self.assertEqual(
+            self.get_postgres_column_type(
+                tablename="band", column_name="popularity"
+            ),
+            "BIGINT",
+        )
 
         popularity = (
             Band.select(Band.popularity).first().run_sync()["popularity"]
@@ -177,21 +173,35 @@ def test_integer_to_varchar(self):
         )
         alter_query.run_sync()
 
-        query = """
-            SELECT data_type FROM information_schema.columns
-            WHERE table_name = 'band'
-            AND table_catalog = 'piccolo'
-            AND column_name = 'popularity'
-            """
-
-        response = Band.raw(query).run_sync()
-        self.assertEqual(response[0]["data_type"].upper(), "CHARACTER VARYING")
+        self.assertEqual(
+            self.get_postgres_column_type(
+                tablename="band", column_name="popularity"
+            ),
+            "CHARACTER VARYING",
+        )
 
         popularity = (
             Band.select(Band.popularity).first().run_sync()["popularity"]
         )
         self.assertEqual(popularity, "1000")
 
+    def test_using_expression(self):
+        """
+        Test the `using_expression` option, which can be used to tell Postgres
+        how to convert certain column types.
+        """
+        Band(name="1").save().run_sync()
+
+        alter_query = Band.alter().set_column_type(
+            old_column=Band.name,
+            new_column=Integer(),
+            using_expression="name::integer",
+        )
+        alter_query.run_sync()
+
+        popularity = Band.select(Band.name).first().run_sync()["name"]
+        self.assertEqual(popularity, 1)
+
 
 @postgres_only
 class TestSetNull(DBTestCase):

Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,7 @@ def __sub__(self, value: DiffableTable) -> TableDelta:`
`102`	`102`	`for key, _ in delta.items()`
`103`	`103`	`}`
`104`	`104`
`105`		`- if delta:`
	`105`	`+ if delta or (column.__class__ != existing_column.__class__):`
`106`	`106`	`alter_columns.append(`
`107`	`107`	`AlterColumn(`
`108`	`108`	`table_class_name=self.class_name,`