From 4c35c44bd2f8469926551681bf604ebd02ef6e30 Mon Sep 17 00:00:00 2001
From: Isaque Alves <isaquealves@gmail.com>
Date: Wed, 22 Jun 2022 09:09:53 -0300
Subject: [PATCH] feat: Add support for dynamically load DDL classes

Adopt a strategy of loading classes based on their names, allowing to
easily add new database support without changing Migrate class logic
---
 aerich/migrate.py | 20 ++++++++------------
 1 file changed, 8 insertions(+), 12 deletions(-)

diff --git a/aerich/migrate.py b/aerich/migrate.py
index 6085efc..e114858 100644
--- a/aerich/migrate.py
+++ b/aerich/migrate.py
@@ -1,3 +1,4 @@
+import importlib
 import os
 from datetime import datetime
 from hashlib import md5
@@ -63,6 +64,11 @@ class Migrate:
             ret = await connection.execute_query(sql)
             cls._db_version = ret[1][0].get("version")
 
+    @classmethod
+    async def load_ddl_class(cls):
+        ddl_dialect_module = importlib.import_module(f"aerich.ddl.{cls.dialect}")
+        return getattr(ddl_dialect_module, f"{cls.dialect.capitalize()}DDL")
+
     @classmethod
     async def init(cls, config: dict, app: str, location: str):
         await Tortoise.init(config=config)
@@ -74,18 +80,8 @@ class Migrate:
 
         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)
+        cls.ddl_class = await cls.load_ddl_class()
+        cls.ddl = cls.ddl_class(connection)
         await cls._get_db_version(connection)
 
     @classmethod