def setUp(self) -> None:
client = Tortoise.get_connection("models")
if client.schema_generator is MySQLSchemaGenerator:
self.ddl = MysqlDDL(client)
elif client.schema_generator is SqliteSchemaGenerator:
self.ddl = SqliteDDL(client)
elif client.schema_generator is AsyncpgSchemaGenerator:
self.ddl = PostgresDDL(client)
async def init_with_old_models(cls, config: dict, app: str, location: str):
await Tortoise.init(config=config)
last_version = await cls.get_last_version()
cls.app = app
cls.migrate_location = os.path.join(location, app)
if last_version:
content = last_version.content
with open(cls.get_old_model_file(app, location), "w", encoding="utf-8") as f:
f.write(content)
migrate_config = cls._get_migrate_config(config, app, location)
cls.migrate_config = migrate_config
await Tortoise.init(config=migrate_config)
connection = get_app_connection(config, app)
cls.dialect = connection.schema_generator.DIALECT
if cls.dialect == "mysql":
from aerich.ddl.mysql import MysqlDDL
cls.ddl = MysqlDDL(connection)
elif cls.dialect == "sqlite":
from aerich.ddl.sqlite import SqliteDDL
cls.ddl = SqliteDDL(connection)
elif cls.dialect == "postgres":
from aerich.ddl.postgres import PostgresDDL
cls.ddl = PostgresDDL(connection)
async def initialize_tests(event_loop, request):
await Tortoise.init(config=tortoise_orm, _create_db=True)
await generate_schema_for_client(Tortoise.get_connection("default"),
safe=True)
client = Tortoise.get_connection("default")
if client.schema_generator is MySQLSchemaGenerator:
Migrate.ddl = MysqlDDL(client)
elif client.schema_generator is SqliteSchemaGenerator:
Migrate.ddl = SqliteDDL(client)
elif client.schema_generator is AsyncpgSchemaGenerator:
Migrate.ddl = PostgresDDL(client)
Migrate.dialect = Migrate.ddl.DIALECT
request.addfinalizer(
lambda: event_loop.run_until_complete(Tortoise._drop_databases()))
async def initialize_tests(event_loop, request):
tortoise_orm["connections"]["diff_models"] = "sqlite://:memory:"
tortoise_orm["apps"]["diff_models"] = {
"models": ["tests.diff_models"],
"default_connection": "diff_models",
}
await Tortoise.init(config=tortoise_orm, _create_db=True)
await generate_schema_for_client(Tortoise.get_connection("default"), safe=True)
client = Tortoise.get_connection("default")
if client.schema_generator is MySQLSchemaGenerator:
Migrate.ddl = MysqlDDL(client)
elif client.schema_generator is SqliteSchemaGenerator:
Migrate.ddl = SqliteDDL(client)
elif client.schema_generator is AsyncpgSchemaGenerator:
Migrate.ddl = PostgresDDL(client)
request.addfinalizer(lambda: event_loop.run_until_complete(Tortoise._drop_databases()))
async def init_with_old_models(cls, config: dict, app: str, location: str):
migrate_config = cls._get_migrate_config(config, app, location)
cls.app = app
cls.migrate_config = migrate_config
cls.migrate_location = os.path.join(location, app)
await Tortoise.init(config=migrate_config)
connection = get_app_connection(config, app)
if connection.schema_generator.DIALECT == "mysql":
from aerich.ddl.mysql import MysqlDDL
cls.ddl = MysqlDDL(connection)
elif connection.schema_generator.DIALECT == "sqlite":
from aerich.ddl.sqlite import SqliteDDL
cls.ddl = SqliteDDL(connection)
elif connection.schema_generator.DIALECT == "postgres":
from aerich.ddl.postgres import PostgresDDL
cls.ddl = PostgresDDL(connection)
async def init(cls, config: dict, app: str, location: str):
await Tortoise.init(config=config)
last_version = await cls.get_last_version()
cls.app = app
cls.migrate_location = Path(location, app)
if last_version:
cls._last_version_content = last_version.content
connection = get_app_connection(config, app)
cls.dialect = connection.schema_generator.DIALECT
if cls.dialect == "mysql":
from aerich.ddl.mysql import MysqlDDL
cls.ddl = MysqlDDL(connection)
elif cls.dialect == "sqlite":
from aerich.ddl.sqlite import SqliteDDL
cls.ddl = SqliteDDL(connection)
elif cls.dialect == "postgres":
from aerich.ddl.postgres import PostgresDDL
cls.ddl = PostgresDDL(connection)
await cls._get_db_version(connection)
class TestDDL(test.TruncationTestCase):
maxDiff = None
def setUp(self) -> None:
client = Tortoise.get_connection("models")
if client.schema_generator is MySQLSchemaGenerator:
self.ddl = MysqlDDL(client)
elif client.schema_generator is SqliteSchemaGenerator:
self.ddl = SqliteDDL(client)
elif client.schema_generator is AsyncpgSchemaGenerator:
self.ddl = PostgresDDL(client)
def test_create_table(self):
ret = self.ddl.create_table(Category)
if isinstance(self.ddl, MysqlDDL):
self.assertEqual(
ret,
"""CREATE TABLE IF NOT EXISTS `category` (
`id` INT NOT NULL PRIMARY KEY AUTO_INCREMENT,
`slug` VARCHAR(200) NOT NULL,
`name` VARCHAR(200) NOT NULL,
`created_at` DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
`user_id` INT NOT NULL COMMENT 'User',
CONSTRAINT `fk_category_user_e2e3874c` FOREIGN KEY (`user_id`) REFERENCES `user` (`id`) ON DELETE CASCADE
) CHARACTER SET utf8mb4;""",
)
elif isinstance(self.ddl, SqliteDDL):
self.assertEqual(
ret,
"""CREATE TABLE IF NOT EXISTS "category" (
"id" INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
"slug" VARCHAR(200) NOT NULL,
"name" VARCHAR(200) NOT NULL,
"created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
"user_id" INT NOT NULL REFERENCES "user" ("id") ON DELETE CASCADE /* User */
);""",
)
elif isinstance(self.ddl, PostgresDDL):
self.assertEqual(
ret,
"""CREATE TABLE IF NOT EXISTS "category" (
"id" SERIAL NOT NULL PRIMARY KEY,
"slug" VARCHAR(200) NOT NULL,
"name" VARCHAR(200) NOT NULL,
"created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
"user_id" INT NOT NULL REFERENCES "user" ("id") ON DELETE CASCADE
);
COMMENT ON COLUMN "category"."user_id" IS 'User';""",
)
def test_drop_table(self):
ret = self.ddl.drop_table(Category)
self.assertEqual(ret, "DROP TABLE IF EXISTS category")
def test_add_column(self):
ret = self.ddl.add_column(Category,
Category._meta.fields_map.get("name"))
if isinstance(self.ddl, MysqlDDL):
self.assertEqual(
ret, "ALTER TABLE category ADD `name` VARCHAR(200) NOT NULL")
elif isinstance(self.ddl, PostgresDDL):
self.assertEqual(
ret, 'ALTER TABLE category ADD "name" VARCHAR(200) NOT NULL')
elif isinstance(self.ddl, SqliteDDL):
self.assertEqual(
ret, 'ALTER TABLE category ADD "name" VARCHAR(200) NOT NULL')
def test_drop_column(self):
ret = self.ddl.drop_column(Category, "name")
self.assertEqual(ret, "ALTER TABLE category DROP COLUMN name")
self.assertEqual(ret, "ALTER TABLE category DROP COLUMN name")
def test_add_index(self):
index = self.ddl.add_index(Category, ["name"])
index_u = self.ddl.add_index(Category, ["name"], True)
if isinstance(self.ddl, MysqlDDL):
self.assertEqual(
index,
"ALTER TABLE category ADD INDEX idx_category_name_8b0cb9 (`name`)"
)
self.assertEqual(
index_u,
"ALTER TABLE category ADD UNIQUE INDEX uid_category_name_8b0cb9 (`name`)"
)
elif isinstance(self.ddl, SqliteDDL):
self.assertEqual(
index_u,
'ALTER TABLE category ADD UNIQUE INDEX uid_category_name_8b0cb9 ("name")'
)
self.assertEqual(
index_u,
'ALTER TABLE category ADD UNIQUE INDEX uid_category_name_8b0cb9 ("name")'
)
def test_drop_index(self):
ret = self.ddl.drop_index(Category, ["name"])
self.assertEqual(
ret, "ALTER TABLE category DROP INDEX idx_category_name_8b0cb9")
ret = self.ddl.drop_index(Category, ["name"], True)
self.assertEqual(
ret, "ALTER TABLE category DROP INDEX uid_category_name_8b0cb9")
def test_add_fk(self):
ret = self.ddl.add_fk(Category, Category._meta.fields_map.get("user"))
self.assertEqual(
ret,
"ALTER TABLE category ADD CONSTRAINT `fk_category_user_e2e3874c` FOREIGN KEY (`user_id`) REFERENCES `user` (`id`) ON DELETE CASCADE",
)
def test_drop_fk(self):
ret = self.ddl.drop_fk(Category, Category._meta.fields_map.get("user"))
self.assertEqual(
ret,
"ALTER TABLE category DROP FOREIGN KEY fk_category_user_e2e3874c")