Make the multi-database configuration fully automatic

docs/index.rst

@@ -65,6 +65,17 @@ To see all the commands that are available run this command::
 
     $ python app.py db --help
 
+Multiple Database Support
+-------------------------
+
+Flask-Migrate can integrate with the  `binds <https://pythonhosted.org/Flask-SQLAlchemy/binds.html>` feature of Flask-SQLAlchemy, making it possible to track migrations to multiple databases associated with an application.
+
+To create a multiple database migration repository, add the ``--multidb`` argument to the ``init`` command::
+
+    $ python app.py db init --multidb
+
+With this command, the migration repository will be set up to track migrations on your main database, and on any additional databases defined in the ``SQLALCHEMY_BINDS`` configuration option.
+
 Command Reference
 -----------------
 
@@ -83,20 +94,18 @@ The application will now have a ``db`` command line option with several sub-comm
     Shows a list of available commands.
 
 - ``manage.py db init [--multidb]``
-    Initializes migration support for the application. Turning on option ``--multidb`` will create multiple databases templates for alembic. This feature could be used with `Flask-SQLAlchemy Binds<https://pythonhosted.org/Flask-SQLAlchemy/binds.html>`. Note that you do *NOT* need this option for other commands, e.g. migrate, upgrade, downgrade, etc. Two more steps are needed once you get the alembic template files:
-    1. Add all database names to the filed ``databases`` in ``alembic.ini``. The ``SQLALCHEMY_DATABASE_URI`` is by default already set as "primary" (you can customize it, but make sure it also gets updated in the ``env.py``), all keys in the ``SQLALCHEMY_BINDS`` should be append to ``database`` filed as a comma seperated string.
-    2. Set ``target_metadata`` in ``env.py``, each database should have a ``db`` object, which, in turn, has all the table information in the ``metadata``. See more detail in the template comment.
+    Initializes migration support for the application. The optional ``--multidb`` enables migrations for multiple databases, configured as `Flask-SQLAlchemy binds <https://pythonhosted.org/Flask-SQLAlchemy/binds.html>`.
 
 - ``manage.py db revision [--message MESSAGE] [--autogenerate] [--sql] [--head HEAD] [--splice] [--branch-label BRANCH_LABEL] [--version-path VERSION_PATH] [--rev-id REV_ID]``
     Creates an empty revision script. The script needs to be edited manually with the upgrade and downgrade changes. See `Alembic's documentation <https://alembic.readthedocs.org/en/latest/index.html>`_ for instructions on how to write migration scripts. An optional migration message can be included.
 
 - ``manage.py db migrate [--message MESSAGE] [--sql] [--head HEAD] [--splice] [--branch-label BRANCH_LABEL] [--version-path VERSION_PATH] [--rev-id REV_ID]``
     Equivalent to ``revision --autogenerate``. The migration script is populated with changes detected automatically. The generated script should to be reviewed and edited as not all types of changes can be detected. This command does not make any changes to the database.
 
-- ``manage.py db upgrade [--sql] [--tag TAG] <revision>``
+- ``manage.py db upgrade [--sql] [--tag TAG] [--x-arg ARG] <revision>``
     Upgrades the database. If ``revision`` isn't given then ``"head"`` is assumed.
 
-- ``manage.py db downgrade [--sql] [--tag TAG] <revision>``
+- ``manage.py db downgrade [--sql] [--tag TAG] [--x-arg ARG] <revision>``
     Downgrades the database. If ``revision`` isn't given then ``-1`` is assumed.
 
 - ``manage.py db stamp [--sql] [--tag TAG] <revision>``

flask_migrate/__init__.py

@@ -59,9 +59,9 @@ def _get_config(directory, x_arg=None):
 @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
                        help=("migration script directory (default is "
                              "'migrations')"))
-@MigrateCommand.option('-m', '--multidb', dest='multidb', action='store_true',
+@MigrateCommand.option('--multidb', dest='multidb', action='store_true',
                        default=False,
-                       help=("multiple databases migraton (default is "
+                       help=("Multiple databases migraton (default is "
                              "False)"))
 def init(directory=None, multidb=False):
     """Generates a new migration"""

flask_migrate/templates/flask-multidb/alembic.ini.mako

@@ -8,11 +8,6 @@
 # the 'revision' command, regardless of autogenerate
 # revision_environment = false
 
-# comma seperated database names, the default database
-# (FLASK_SQLALCHEMY_URL) is primary, other names must be
-# the same as names in SQLALCHEMY_BINDS
-# e.g. database = primary, db1, db2, ...
-databases = primary
 
 # Logging configuration
 [loggers]

flask_migrate/templates/flask-multidb/env.py

@@ -1,9 +1,8 @@
 from __future__ import with_statement
 from alembic import context
-from sqlalchemy import engine_from_config, pool
+from sqlalchemy import engine_from_config, pool, MetaData
 from logging.config import fileConfig
 import logging
-import re
 
 USE_TWOPHASE = False
 
@@ -16,39 +15,37 @@
 fileConfig(config.config_file_name)
 logger = logging.getLogger('alembic.env')
 
-# gather section names referring to different
-# databases.
-db_names = config.get_main_option('databases')
-
-# gather the database engine's information
+# add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
 from flask import current_app
 config.set_main_option('sqlalchemy.url',
                        current_app.config.get('SQLALCHEMY_DATABASE_URI'))
-context.config.set_section_option("primary", "sqlalchemy.url",
-                                  current_app.config.get('SQLALCHEMY_DATABASE_URI'))
-for engine, url in current_app.config.get("SQLALCHEMY_BINDS").items():
-    context.config.set_section_option(engine, "sqlalchemy.url", url)
-
-# add your model's MetaData objects here
-# for 'autogenerate' support.  These must be set
-# up to hold just those tables targeting a
-# particular database. table.tometadata() may be
-# helpful here in case a "copy" of
-# a MetaData is needed.
-# from myapp import mymodel
-# target_metadata = {
-#       'engine1':mymodel.metadata1,
-#       'engine2':mymodel.metadata2
-#}
-target_metadata = {
-}
+bind_names = []
+for name, url in current_app.config.get("SQLALCHEMY_BINDS").items():
+    context.config.set_section_option(name, "sqlalchemy.url", url)
+    bind_names.append(name)
+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.
 
@@ -64,19 +61,19 @@ def run_migrations_offline():
     # for the --sql use case, run migrations for each URL into
     # individual files.
 
-    engines = {}
-    for name in re.split(r',\s*', db_names):
+    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)
+        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=target_metadata.get(name),
+                              target_metadata=get_metadata(name),
                               literal_binds=True)
             with context.begin_transaction():
                 context.run_migrations(engine_name=name)
@@ -93,8 +90,11 @@ def run_migrations_online():
     # for the direct-to-DB use case, start a transaction on all
     # engines, then run all migrations, then commit all transactions.
 
-    engines = {}
-    for name in re.split(r',\s*', db_names):
+    engines = {'': {'engine': engine_from_config(
+        config.get_section(config.config_ini_section),
+        prefix='sqlalchemy.',
+        poolclass=pool.NullPool)}}
+    for name in bind_names:
         engines[name] = rec = {}
         rec['engine'] = engine_from_config(
             context.config.get_section(name),
@@ -112,12 +112,12 @@ def run_migrations_online():
 
     try:
         for name, rec in engines.items():
-            logger.info("Migrating database %s" % name)
+            logger.info("Migrating database %s" % (name or '<default>'))
             context.configure(
                 connection=rec['connection'],
                 upgrade_token="%s_upgrades" % name,
                 downgrade_token="%s_downgrades" % name,
-                target_metadata=target_metadata.get(name)
+                target_metadata=get_metadata(name)
             )
             context.run_migrations(engine_name=name)
 

flask_migrate/templates/flask-multidb/script.py.mako

@@ -27,13 +27,14 @@ def downgrade(engine_name):
     globals()["downgrade_%s" % engine_name]()
 
 <%
-    db_names = config.get_main_option("databases")
+    from flask import current_app
+    db_names = [''] + list(current_app.config.get("SQLALCHEMY_BINDS").keys())
 %>
 
 ## generate an "upgrade_<xyz>() / downgrade_<xyz>()" function
 ## for each database name in the ini file.
 
-% for db_name in re.split(r',\s*', db_names):
+% for db_name in db_names:
 
 def upgrade_${db_name}():
     ${context.get("%s_upgrades" % db_name, "pass")}

tests/multidb/app_multidb.py → tests/app_multidb.py

@@ -4,30 +4,28 @@
 from flask_migrate import Migrate, MigrateCommand
 
 app = Flask(__name__)
-app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///app.db'
+app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///app1.db'
 app.config['SQLALCHEMY_BINDS'] = {
-    "db1": "sqlite:///app1.db",
+    "db1": "sqlite:///app2.db",
 }
 
 db = SQLAlchemy(app)
-migrate = Migrate(app, db)
-db1 = SQLAlchemy(app)
 
-manager = Manager(app)
-manager.add_command('db', MigrateCommand)
 
-metadata = db.metadata
-metadata1 = db1.metadata
+class User(db.Model):
+    id = db.Column(db.Integer, primary_key=True)
+    name = db.Column(db.String(128))
 
 
-class User(db.Model):
+class Group(db.Model):
+    __bind_key__ = 'db1'
     id = db.Column(db.Integer, primary_key=True)
     name = db.Column(db.String(128))
 
+migrate = Migrate(app, db)
 
-class Group(db1.Model):
-    id = db1.Column(db1.Integer, primary_key=True)
-    name = db1.Column(db1.String(128))
+manager = Manager(app)
+manager.add_command('db', MigrateCommand)
 
 
 if __name__ == '__main__':