Better handling of MetaData instances in templates

Author: miguelgrinberg

src/flask_migrate/templates/aioflask-multidb/env.py

@@ -41,7 +41,7 @@
         bind, "sqlalchemy.url",
         str(current_app.extensions['migrate'].db.get_engine(
             bind=bind).url).replace('%', '%%'))
-target_metadata = current_app.extensions['migrate'].db.metadata
+target_db = current_app.extensions['migrate'].db
 
 
 # other values from the config, defined by the needs of env.py,
@@ -54,8 +54,12 @@ def get_metadata(bind):
     """Return the metadata for a bind."""
     if bind == '':
         bind = None
+    if hasattr(target_db, 'metadatas'):
+        return target_db.metadatas[bind]
+
+    # legacy, less flexible implementation
     m = MetaData()
-    for t in target_metadata.tables.values():
+    for t in target_db.metadata.tables.values():
         if t.info.get('bind_key') == bind:
             t.tometadata(m)
     return m
@@ -182,4 +186,4 @@ async def run_migrations_online():
 if context.is_offline_mode():
     run_migrations_offline()
 else:
-    asyncio.run(run_migrations_online())
+    asyncio.get_event_loop().run_until_complete(run_migrations_online())

src/flask_migrate/templates/aioflask/env.py

@@ -25,14 +25,20 @@
     'sqlalchemy.url',
     str(current_app.extensions['migrate'].db.get_engine().url).replace(
         '%', '%%'))
-target_metadata = current_app.extensions['migrate'].db.metadata
+target_db = current_app.extensions['migrate'].db
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
 # my_important_option = config.get_main_option("my_important_option")
 # ... etc.
 
 
+def get_metadata():
+    if hasattr(target_db, 'metadatas'):
+        return target_db.metadatas[None]
+    return target_db.metadata
+
+
 def run_migrations_offline():
     """Run migrations in 'offline' mode.
 
@@ -47,7 +53,7 @@ def run_migrations_offline():
     """
     url = config.get_main_option("sqlalchemy.url")
     context.configure(
-        url=url, target_metadata=target_metadata, literal_binds=True
+        url=url, target_metadata=get_metadata(), literal_binds=True
     )
 
     with context.begin_transaction():
@@ -67,7 +73,7 @@ def process_revision_directives(context, revision, directives):
 
     context.configure(
         connection=connection,
-        target_metadata=target_metadata,
+        target_metadata=get_metadata(),
         process_revision_directives=process_revision_directives,
         **current_app.extensions['migrate'].configure_args
     )
@@ -93,4 +99,4 @@ async def run_migrations_online():
 if context.is_offline_mode():
     run_migrations_offline()
 else:
-    asyncio.run(run_migrations_online())
+    asyncio.get_event_loop().run_until_complete(run_migrations_online())

src/flask_migrate/templates/flask-multidb/env.py

@@ -40,8 +40,7 @@
         bind, "sqlalchemy.url",
         str(current_app.extensions['migrate'].db.get_engine(
             bind=bind).url).replace('%', '%%'))
-target_metadata = current_app.extensions['migrate'].db.metadata
-
+target_db = current_app.extensions['migrate'].db
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
@@ -53,8 +52,12 @@ def get_metadata(bind):
     """Return the metadata for a bind."""
     if bind == '':
         bind = None
+    if hasattr(target_db, 'metadatas'):
+        return target_db.metadatas[bind]
+
+    # legacy, less flexible implementation
     m = MetaData()
-    for t in target_metadata.tables.values():
+    for t in target_db.metadata.tables.values():
         if t.info.get('bind_key') == bind:
             t.tometadata(m)
     return m

src/flask_migrate/templates/flask/env.py

@@ -24,14 +24,20 @@
     'sqlalchemy.url',
     str(current_app.extensions['migrate'].db.get_engine().url).replace(
         '%', '%%'))
-target_metadata = current_app.extensions['migrate'].db.metadata
+target_db = current_app.extensions['migrate'].db
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
 # my_important_option = config.get_main_option("my_important_option")
 # ... etc.
 
 
+def get_metadata():
+    if hasattr(target_db, 'metadatas'):
+        return target_db.metadatas[None]
+    return target_db.metadata
+
+
 def run_migrations_offline():
     """Run migrations in 'offline' mode.
 
@@ -46,7 +52,7 @@ def run_migrations_offline():
     """
     url = config.get_main_option("sqlalchemy.url")
     context.configure(
-        url=url, target_metadata=target_metadata, literal_binds=True
+        url=url, target_metadata=get_metadata(), literal_binds=True
     )
 
     with context.begin_transaction():
@@ -76,7 +82,7 @@ def process_revision_directives(context, revision, directives):
     with connectable.connect() as connection:
         context.configure(
             connection=connection,
-            target_metadata=target_metadata,
+            target_metadata=get_metadata(),
             process_revision_directives=process_revision_directives,
             **current_app.extensions['migrate'].configure_args
         )