def test_alter_column(self):
"""
Test altering columns.
"""
manager_1 = MigrationManager()
manager_1.add_table(class_name="Manager", tablename="manager")
manager_1.add_column(
table_class_name="Manager",
tablename="manager",
column_name="name",
column_class_name="Varchar",
params={"length": 100},
)
manager_2 = MigrationManager()
manager_2.alter_column(
table_class_name="Manager",
tablename="manager",
column_name="name",
params={"unique": True},
old_params={"unique": False},
)
schema_snapshot = SchemaSnapshot(managers=[manager_1, manager_2])
snapshot = schema_snapshot.get_snapshot()
self.assertTrue(snapshot[0].columns[0]._meta.unique)
async def forwards():
manager = MigrationManager(migration_id=ID,
app_name="music",
description=DESCRIPTION)
manager.add_column(
table_class_name="Ticket",
tablename="ticket",
column_name="concert",
column_class_name="ForeignKey",
column_class=ForeignKey,
params={
"references": Concert,
"on_delete": OnDelete.cascade,
"on_update": OnUpdate.cascade,
"null": True,
"primary_key": False,
"unique": False,
"index": False,
"index_method": IndexMethod.btree,
"choices": None,
},
)
return manager
def test_add_non_nullable_column(self):
"""
Test adding a non nullable column to a MigrationManager.
Need to handle it gracefully if rows already exist.
"""
self.run_sync("INSERT INTO manager VALUES (default, 'Dave');")
manager = MigrationManager()
manager.add_column(
table_class_name="Manager",
tablename="manager",
column_name="email",
column_class_name="Varchar",
params={
"length": 100,
"default": "",
"null": False,
"primary": False,
"key": False,
"unique": True,
"index": False,
},
)
asyncio.run(manager.run())
async def forwards():
manager = MigrationManager(migration_id=ID,
app_name="music",
description=DESCRIPTION)
manager.alter_column(
table_class_name="Venue",
tablename="venue",
column_name="capacity",
params={"secret": True},
old_params={"secret": False},
column_class=Integer,
old_column_class=Integer,
)
return manager
def test_drop_table(self):
"""
Test dropping tables.
"""
manager_1 = MigrationManager()
manager_1.add_table(class_name="Manager", tablename="manager")
manager_1.add_table(class_name="Band", tablename="band")
manager_2 = MigrationManager()
manager_2.drop_table(class_name="Band", tablename="band")
schema_snapshot = SchemaSnapshot(managers=[manager_1, manager_2])
snapshot = schema_snapshot.get_snapshot()
self.assertTrue(len(snapshot) == 1)
class_names = [i.class_name for i in snapshot]
self.assertTrue("Manager" in class_names)
async def forwards():
manager = MigrationManager(migration_id=ID, app_name="questions")
manager.add_column(
table_class_name="Category",
tablename="category",
column_name="slug",
column_class_name="Varchar",
params={
"length": 200,
"default": "",
"null": False,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
return manager
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,
)
def test_alter_column_set_default(self):
"""
Test altering a column default with MigrationManager.
"""
manager = MigrationManager()
manager.alter_column(
table_class_name="Manager",
tablename="manager",
column_name="name",
params={"default": "Unknown"},
old_params={"default": ""},
)
asyncio.run(manager.run())
self.assertEqual(
self._get_column_default(),
[{"column_default": "'Unknown'::character varying"}],
)
asyncio.run(manager.run_backwards())
self.assertEqual(
self._get_column_default(),
[{"column_default": "''::character varying"}],
)
def test_add_column(self):
"""
Test adding a column 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": False,
},
)
asyncio.run(manager.run())
self.run_sync(
"INSERT INTO manager VALUES (default, 'Dave', '*****@*****.**');"
)
response = self.run_sync("SELECT * FROM manager;")
self.assertEqual(
response, [{"id": 1, "name": "Dave", "email": "*****@*****.**"}]
)
# Reverse
asyncio.run(manager.run_backwards())
response = self.run_sync("SELECT * FROM manager;")
self.assertEqual(response, [{"id": 1, "name": "Dave"}])
def test_alter_column_unique(self):
"""
Test altering a column uniqueness with MigrationManager.
"""
manager = MigrationManager()
manager.alter_column(
table_class_name="Manager",
tablename="manager",
column_name="name",
params={"unique": True},
old_params={"unique": False},
)
asyncio.run(manager.run())
with self.assertRaises(UniqueViolationError):
self.run_sync(
"INSERT INTO manager VALUES "
"(default, 'Dave'), (default, 'Dave');"
)
# Reverse
asyncio.run(manager.run_backwards())
self.run_sync(
"INSERT INTO manager VALUES (default, 'Dave'), (default, 'Dave');"
)
response = self.run_sync("SELECT name FROM manager;")
self.assertEqual(response, [{"name": "Dave"}, {"name": "Dave"}])
def test_rename_column(self):
"""
Test running a MigrationManager which contains a column rename
operation.
"""
self.insert_row()
manager = MigrationManager()
manager.rename_column(
table_class_name="Band",
tablename="band",
old_column_name="name",
new_column_name="title",
)
asyncio.run(manager.run())
response = self.run_sync("SELECT * FROM band;")
self.assertTrue("title" in response[0].keys())
self.assertTrue("name" not in response[0].keys())
# Reverse
asyncio.run(manager.run_backwards())
response = self.run_sync("SELECT * FROM band;")
self.assertTrue("title" not in response[0].keys())
self.assertTrue("name" in response[0].keys())
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")
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 test_alter_column_digits(self):
"""
Test altering a column digits with MigrationManager.
"""
manager = MigrationManager()
manager.alter_column(
table_class_name="Ticket",
tablename="ticket",
column_name="price",
params={"digits": (6, 2)},
old_params={"digits": (5, 2)},
)
asyncio.run(manager.run())
self.assertEqual(
self._get_column_precision_and_scale(),
[{"numeric_precision": 6, "numeric_scale": 2}],
)
asyncio.run(manager.run_backwards())
self.assertEqual(
self._get_column_precision_and_scale(),
[{"numeric_precision": 5, "numeric_scale": 2}],
)
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())
def test_drop_column(self):
"""
Test dropping columns.
"""
manager_1 = MigrationManager()
manager_1.add_table(class_name="Manager", tablename="manager")
manager_1.add_column(
table_class_name="Manager",
tablename="manager",
column_name="name",
column_class_name="Varchar",
params={"length": 100},
)
manager_2 = MigrationManager()
manager_2.drop_column(
table_class_name="Manager", tablename="manager", column_name="name"
)
schema_snapshot = SchemaSnapshot(managers=[manager_1, manager_2])
snapshot = schema_snapshot.get_snapshot()
self.assertEqual(len(snapshot[0].columns), 0)
def test_get_table_from_snapshot(self):
manager_1 = MigrationManager()
manager_1.add_table(class_name="Manager", tablename="manager")
manager_1.add_table(class_name="Band", tablename="band")
schema_snapshot = SchemaSnapshot(managers=[manager_1])
table = schema_snapshot.get_table_from_snapshot("Manager")
self.assertTrue(table.class_name == "Manager")
with self.assertRaises(ValueError):
schema_snapshot.get_table_from_snapshot("Foo")
def test_raw_coroutine(self):
"""
Test adding raw coroutines to a MigrationManager.
"""
class HasRun(Exception):
pass
async def run():
raise HasRun("I was run!")
manager = MigrationManager()
manager.add_raw(run)
manager.add_raw_backwards(run)
with self.assertRaises(HasRun):
asyncio.run(manager.run())
# Reverse
with self.assertRaises(HasRun):
asyncio.run(manager.run_backwards())
def test_rename_table(self):
"""
Test renaming tables.
"""
manager_1 = MigrationManager()
manager_1.add_table(class_name="Band", tablename="band")
manager_2 = MigrationManager()
manager_2.rename_table(
old_class_name="Band",
old_tablename="band",
new_class_name="Performer",
new_tablename="performer",
)
schema_snapshot = SchemaSnapshot(managers=[manager_1, manager_2])
snapshot = schema_snapshot.get_snapshot()
self.assertTrue(snapshot[0].class_name == "Performer")
self.assertTrue(snapshot[0].tablename == "performer")
async def forwards():
manager = MigrationManager(migration_id=ID, app_name="token_auth")
manager.add_table("TokenAuth", tablename="token_auth")
manager.add_column(
table_class_name="TokenAuth",
tablename="token_auth",
column_name="token",
column_class_name="Varchar",
params={
"length": 255,
"default": generate_token,
"null": False,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
manager.add_column(
table_class_name="TokenAuth",
tablename="token_auth",
column_name="user",
column_class_name="ForeignKey",
params={
"references": BaseUser,
"on_delete": OnDelete.cascade,
"on_update": OnUpdate.cascade,
"default": None,
"null": True,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
return manager
async def forwards():
manager = MigrationManager(migration_id=ID, app_name="forum")
manager.drop_column(table_class_name="Topic",
tablename="topic",
column_name="description")
manager.rename_column(
table_class_name="Topic",
tablename="topic",
old_column_name="title",
new_column_name="subject",
)
return manager
async def forwards():
manager = MigrationManager(migration_id=ID, app_name="forum")
manager.add_table("Category", tablename="category")
manager.add_column(
table_class_name="Category",
tablename="category",
column_name="name",
column_class_name="Varchar",
params={
"length": 255,
"default": "",
"null": False,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
manager.add_column(
table_class_name="Category",
tablename="category",
column_name="description",
column_class_name="Text",
params={
"default": "",
"null": False,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
return manager
async def forwards():
manager = MigrationManager(migration_id=ID, app_name="home")
manager.add_table("Task", tablename="task")
manager.add_column(
table_class_name="Task",
tablename="task",
column_name="name",
column_class_name="Varchar",
params={
"length": 255,
"default": "",
"null": False,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
manager.add_column(
table_class_name="Task",
tablename="task",
column_name="completed",
column_class_name="Boolean",
params={
"default": False,
"null": False,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
return manager
def test_add_foreign_key_self_column(self):
"""
Test adding a ForeignKey column to a MigrationManager, with a
references argument of 'self'.
"""
manager = MigrationManager()
manager.add_column(
table_class_name="Manager",
tablename="manager",
column_name="advisor",
column_class=ForeignKey,
column_class_name="ForeignKey",
params={
"references": "self",
"on_delete": OnDelete.cascade,
"on_update": OnUpdate.cascade,
"default": None,
"null": True,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
asyncio.run(manager.run())
self.run_sync("INSERT INTO manager VALUES (default, 'Alice', null);")
self.run_sync("INSERT INTO manager VALUES (default, 'Dave', 1);")
response = self.run_sync("SELECT * FROM manager;")
self.assertEqual(
response,
[
{"id": 1, "name": "Alice", "advisor": None},
{"id": 2, "name": "Dave", "advisor": 1},
],
)
# Reverse
asyncio.run(manager.run_backwards())
response = self.run_sync("SELECT * FROM manager;")
self.assertEqual(
response, [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Dave"}],
)
def test_add_column(self):
"""
Test adding columns.
"""
manager = MigrationManager()
manager.add_table(class_name="Manager", tablename="manager")
manager.add_column(
table_class_name="Manager",
tablename="manager",
column_name="name",
column_class_name="Varchar",
params={"length": 100},
)
schema_snapshot = SchemaSnapshot(managers=[manager])
snapshot = schema_snapshot.get_snapshot()
self.assertTrue(len(snapshot) == 1)
self.assertTrue(len(snapshot[0].columns) == 1)
async def forwards():
manager = MigrationManager(migration_id=ID)
manager.add_table("Band", tablename="band")
manager.add_column(
table_class_name="Band",
tablename="band",
column_name="name",
column_class_name="Varchar",
params={
"length": 150,
"default": "",
"null": True,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
return manager
async def forwards():
manager = MigrationManager(migration_id=ID, app_name="music")
manager.add_table("Poster", tablename="poster")
manager.add_column(
table_class_name="Poster",
tablename="poster",
column_name="content",
column_class_name="Text",
params={
"default": "",
"null": False,
"primary": False,
"key": False,
"unique": False,
"index": False,
},
)
return manager
def test_alter_column_add_index(self):
"""
Test altering a column to add an index with MigrationManager.
"""
manager = MigrationManager()
manager.alter_column(
table_class_name="Manager",
tablename="manager",
column_name="name",
params={"index": True},
old_params={"index": False},
)
asyncio.run(manager.run())
self.assertTrue(
Manager._get_index_name(["name"]) in Manager.indexes().run_sync())
asyncio.run(manager.run_backwards())
self.assertTrue(
Manager._get_index_name(["name"]) not in
Manager.indexes().run_sync())
async def forwards():
manager = MigrationManager(migration_id=ID, app_name="user")
manager.add_column(
table_class_name="BaseUser",
tablename="piccolo_user",
column_name="superuser",
column_class_name="Boolean",
column_class=Boolean,
params={
"default": False,
"null": False,
"primary": False,
"key": False,
"unique": False,
"index": False,
"index_method": IndexMethod.btree,
},
)
manager.add_column(
table_class_name="BaseUser",
tablename="piccolo_user",
column_name="last_login",
column_class_name="Timestamp",
column_class=Timestamp,
params={
"default": None,
"null": True,
"primary": False,
"key": False,
"unique": False,
"index": False,
"index_method": IndexMethod.btree,
},
)
return manager
def test_add_table(self):
"""
Test adding a table to a MigrationManager.
"""
self.run_sync("DROP TABLE IF EXISTS musician;")
manager = MigrationManager()
name_column = Varchar()
name_column._meta.name = "name"
manager.add_table(
class_name="Musician", tablename="musician", columns=[name_column]
)
asyncio.run(manager.run())
self.run_sync("INSERT INTO musician VALUES (default, 'Bob Jones');")
response = self.run_sync("SELECT * FROM musician;")
self.assertEqual(response, [{"id": 1, "name": "Bob Jones"}])
# Reverse
asyncio.run(manager.run_backwards())
self.assertEqual(self.table_exists("musician"), False)
self.run_sync("DROP TABLE IF EXISTS musician;")
