Alembic templates for aioflask

src/flask_migrate/templates/aioflask-multidb/README

@@ -0,0 +1 @@
+Multi-database configuration for aioflask and alchemical.

src/flask_migrate/templates/aioflask-multidb/alembic.ini.mako

@@ -0,0 +1,50 @@
+# A generic, single database configuration.
+
+[alembic]
+# template used to generate migration files
+# file_template = %%(rev)s_%%(slug)s
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic,flask_migrate
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[logger_flask_migrate]
+level = INFO
+handlers =
+qualname = flask_migrate
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

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

@@ -0,0 +1,185 @@
+from __future__ import with_statement
+
+import asyncio
+import logging
+from logging.config import fileConfig
+
+from sqlalchemy import MetaData
+from flask import current_app
+
+from alembic import context
+
+USE_TWOPHASE = False
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+fileConfig(config.config_file_name)
+logger = logging.getLogger('alembic.env')
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
+config.set_main_option(
+    'sqlalchemy.url',
+    str(current_app.extensions['migrate'].db.get_engine().url).replace(
+        '%', '%%'))
+bind_names = []
+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(
+            bind=bind).url).replace('%', '%%'))
+target_metadata = current_app.extensions['migrate'].db.metadata
+
+
+# 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(bind):
+    """Return the metadata for a bind."""
+    if bind == '':
+        bind = None
+    m = MetaData()
+    for t in target_metadata.tables.values():
+        if t.info.get('bind_key') == bind:
+            t.tometadata(m)
+    return m
+
+
+def run_migrations_offline():
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    # for the --sql use case, run migrations for each URL into
+    # individual files.
+
+    engines = {
+        '': {
+            'url': context.config.get_main_option('sqlalchemy.url')
+        }
+    }
+    for name in bind_names:
+        engines[name] = rec = {}
+        rec['url'] = context.config.get_section_option(name, "sqlalchemy.url")
+
+    for name, rec in engines.items():
+        logger.info("Migrating database %s" % (name or '<default>'))
+        file_ = "%s.sql" % name
+        logger.info("Writing output to %s" % file_)
+        with open(file_, 'w') as buffer:
+            context.configure(
+                url=rec['url'],
+                output_buffer=buffer,
+                target_metadata=get_metadata(name),
+                literal_binds=True,
+            )
+            with context.begin_transaction():
+                context.run_migrations(engine_name=name)
+
+
+def do_run_migrations(_, engines):
+    # this callback is used to prevent an auto-migration from being generated
+    # when there are no changes to the schema
+    # reference: http://alembic.zzzcomputing.com/en/latest/cookbook.html
+    def process_revision_directives(context, revision, directives):
+        if getattr(config.cmd_opts, 'autogenerate', False):
+            script = directives[0]
+            if len(script.upgrade_ops_list) >= len(bind_names) + 1:
+                empty = True
+                for upgrade_ops in script.upgrade_ops_list:
+                    if not upgrade_ops.is_empty():
+                        empty = False
+                if empty:
+                    directives[:] = []
+                    logger.info('No changes in schema detected.')
+
+    for name, rec in engines.items():
+        rec['sync_connection'] = conn = rec['connection']._sync_connection()
+        if USE_TWOPHASE:
+            rec['transaction'] = conn.begin_twophase()
+        else:
+            rec['transaction'] = conn.begin()
+
+    try:
+        for name, rec in engines.items():
+            logger.info("Migrating database %s" % (name or '<default>'))
+            context.configure(
+                connection=rec['sync_connection'],
+                upgrade_token="%s_upgrades" % name,
+                downgrade_token="%s_downgrades" % name,
+                target_metadata=get_metadata(name),
+                process_revision_directives=process_revision_directives,
+                **current_app.extensions['migrate'].configure_args
+            )
+            context.run_migrations(engine_name=name)
+
+        if USE_TWOPHASE:
+            for rec in engines.values():
+                rec['transaction'].prepare()
+
+        for rec in engines.values():
+            rec['transaction'].commit()
+    except:  # noqa: E722
+        for rec in engines.values():
+            rec['transaction'].rollback()
+        raise
+    finally:
+        for rec in engines.values():
+            rec['sync_connection'].close()
+
+
+async def run_migrations_online():
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+
+    # 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.get_engine()}
+    }
+    for name in bind_names:
+        engines[name] = rec = {}
+        rec['engine'] = current_app.extensions['migrate'].db.get_engine(
+            bind=name)
+
+    for name, rec in engines.items():
+        engine = rec['engine']
+        rec['connection'] = await engine.connect().start()
+
+    await engines['']['connection'].run_sync(do_run_migrations, engines)
+
+    for rec in engines.values():
+        await rec['connection'].close()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    asyncio.run(run_migrations_online())

src/flask_migrate/templates/aioflask-multidb/script.py.mako

@@ -0,0 +1,53 @@
+<%!
+import re
+
+%>"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision = ${repr(up_revision)}
+down_revision = ${repr(down_revision)}
+branch_labels = ${repr(branch_labels)}
+depends_on = ${repr(depends_on)}
+
+
+def upgrade(engine_name):
+    globals()["upgrade_%s" % engine_name]()
+
+
+def downgrade(engine_name):
+    globals()["downgrade_%s" % engine_name]()
+
+<%
+    from flask import current_app
+    bind_names = []
+    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
+## for each database name in the ini file.
+
+% for db_name in db_names:
+
+def upgrade_${db_name}():
+    ${context.get("%s_upgrades" % db_name, "pass")}
+
+
+def downgrade_${db_name}():
+    ${context.get("%s_downgrades" % db_name, "pass")}
+
+% endfor

src/flask_migrate/templates/aioflask/README

@@ -0,0 +1 @@
+Generic single-database configuration with an async engine.

src/flask_migrate/templates/aioflask/alembic.ini.mako

@@ -0,0 +1,50 @@
+# A generic, single database configuration.
+
+[alembic]
+# template used to generate migration files
+# file_template = %%(rev)s_%%(slug)s
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+
+# Logging configuration
+[loggers]
+keys = root,sqlalchemy,alembic,flask_migrate
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[logger_flask_migrate]
+level = INFO
+handlers =
+qualname = flask_migrate
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S