add support for Alchemical in addition to Flask-SQLAlchemy

miguelgrinberg · May 30, 2021 · 113115d · 113115d
1 parent 699e136
commit 113115d
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 10 deletions.
diff --git a/flask_migrate/templates/flask-multidb/env.py b/flask_migrate/templates/flask-multidb/env.py
@@ -25,14 +25,20 @@
 # target_metadata = mymodel.Base.metadata
 config.set_main_option(
     'sqlalchemy.url',
-    str(current_app.extensions['migrate'].db.engine.url).replace('%', '%%'))
-bind_names = []
-for bind in current_app.config.get("SQLALCHEMY_BINDS"):
+    str(current_app.extensions['migrate'].db.get_engine().url).replace(
+        '%', '%%'))
+if current_app.config.get('SQLALCHEMY_BINDS') is not None:
+    bind_names = list(current_app.config['SQLALCHEMY_BINDS'].keys())
+else:
+    get_bind_names = getattr(current_app.extensions['migrate'].db,
+                             'bind_names', None)
+    if get_bind_names:
+        bind_names = get_bind_names()
+for bind in bind_names:
     context.config.set_section_option(
         bind, "sqlalchemy.url",
         str(current_app.extensions['migrate'].db.get_engine(
-            current_app, bind).url).replace('%', '%%'))
-    bind_names.append(bind)
+            bind=bind).url).replace('%', '%%'))
 target_metadata = current_app.extensions['migrate'].db.metadata
 
 
@@ -118,12 +124,12 @@ def process_revision_directives(context, revision, directives):
     # for the direct-to-DB use case, start a transaction on all
     # engines, then run all migrations, then commit all transactions.
     engines = {
-        '': {'engine': current_app.extensions['migrate'].db.engine}
+        '': {'engine': current_app.extensions['migrate'].db.get_engine()}
     }
     for name in bind_names:
         engines[name] = rec = {}
         rec['engine'] = current_app.extensions['migrate'].db.get_engine(
-            app=current_app, bind=name)
+            bind=name)
 
     for name, rec in engines.items():
         engine = rec['engine']

diff --git a/flask_migrate/templates/flask-multidb/script.py.mako b/flask_migrate/templates/flask-multidb/script.py.mako
@@ -28,7 +28,13 @@ def downgrade(engine_name):
 
 <%
     from flask import current_app
-    db_names = [''] + list(current_app.config.get("SQLALCHEMY_BINDS").keys())
+    if current_app.config.get('SQLALCHEMY_BINDS') is not None:
+        bind_names = list(current_app.config['SQLALCHEMY_BINDS'].keys())
+    else:
+        get_bind_names = getattr(current_app.extensions['migrate'].db, 'bind_names', None)
+        if get_bind_names:
+            bind_names = get_bind_names()
+    db_names = [''] + bind_names
 %>
 
 ## generate an "upgrade_<xyz>() / downgrade_<xyz>()" function

diff --git a/flask_migrate/templates/flask/env.py b/flask_migrate/templates/flask/env.py
@@ -22,7 +22,8 @@
 # target_metadata = mymodel.Base.metadata
 config.set_main_option(
     'sqlalchemy.url',
-    str(current_app.extensions['migrate'].db.engine.url).replace('%', '%%'))
+    str(current_app.extensions['migrate'].db.get_engine().url).replace(
+        '%', '%%'))
 target_metadata = current_app.extensions['migrate'].db.metadata
 
 # other values from the config, defined by the needs of env.py,
@@ -70,7 +71,7 @@ def process_revision_directives(context, revision, directives):
                 directives[:] = []
                 logger.info('No changes in schema detected.')
 
-    connectable = current_app.extensions['migrate'].db.engine
+    connectable = current_app.extensions['migrate'].db.get_engine()
 
     with connectable.connect() as connection:
         context.configure(