Merged revision 6ed983fc075450d874557e81feb0237d7a28a222 from upstream default tip

4 files changed, 47 insertions, 24 deletions

diff --git a/alembic/command.py b/alembic/command.py
index ed7b830..67b0d66 100644
--- a/alembic/command.py
+++ b/alembic/command.py
@@ -63,7 +63,9 @@ def revision(config, message=None, autogenerate=False):
     """Create a new revision file."""
 
     script = ScriptDirectory.from_config(config)
-    template_args = {}
+    template_args = {
+        'config': config # Let templates use config for e.g. multiple databases
+    }
     imports = set()
     if autogenerate:
         util.requires_07("autogenerate")
diff --git a/alembic/migration.py b/alembic/migration.py
index 1a65cd2..7669906 100644
--- a/alembic/migration.py
+++ b/alembic/migration.py
@@ -220,7 +220,7 @@ class MigrationContext(object):
             if self.as_sql and not rev:
                 self._version.drop(self.connection)
 
-    def execute(self, sql):
+    def execute(self, sql, execution_options=None):
         """Execute a SQL construct or string statement.
         
         The underlying execution mechanics are used, that is
@@ -229,7 +229,7 @@ class MigrationContext(object):
         the current SQLAlchemy connection.
 
         """
-        self.impl._exec(sql)
+        self.impl._exec(sql, execution_options)
 
     def _stdout_connection(self, connection):
         def dump(construct, *multiparams, **params):
diff --git a/alembic/templates/multidb/env.py b/alembic/templates/multidb/env.py
index 8812294..d3ccc55 100644
--- a/alembic/templates/multidb/env.py
+++ b/alembic/templates/multidb/env.py
@@ -1,17 +1,25 @@
-USE_TWOPHASE = False
-
+from __future__ import with_statement
 from alembic import context
 from sqlalchemy import engine_from_config, pool
+from logging.config import fileConfig
+import logging
 import re
-import sys
 
-import logging
-logging.fileConfig(options.config_file)
+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(__name__)
 
 # gather section names referring to different 
 # databases.  These are named "engine1", "engine2"
 # in the sample .ini file.
-db_names = options.get_main_option('databases')
+db_names = config.get_main_option('databases')
 
 # add your model's MetaData objects here
 # for 'autogenerate' support.  These must be set 
@@ -26,6 +34,11 @@ db_names = options.get_main_option('databases')
 #}
 target_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 run_migrations_offline():
     """Run migrations in 'offline' mode.
 
@@ -48,14 +61,15 @@ def run_migrations_offline():
                                             "sqlalchemy.url")
 
     for name, rec in engines.items():
+        logger.info("Migrating database %s" % name)
         file_ = "%s.sql" % name
-        sys.stderr.write("Writing output to %s\n" % file_)
+        logger.info("Writing output to %s" % file_)
         context.configure(
                     url=rec['url'],
                     output_buffer=open(file_, 'w')
                 )
         with context.begin_transaction():
-            context.run_migrations(engine=name)
+            context.run_migrations(engine_name=name)
 
 def run_migrations_online():
     """Run migrations in 'online' mode.
@@ -87,14 +101,15 @@ def run_migrations_online():
 
     try:
         for name, rec in engines.items():
+            logger.info("Migrating database %s" % name)
             context.configure(
                         connection=rec['connection'],
                         upgrade_token="%s_upgrades",
                         downgrade_token="%s_downgrades",
                         target_metadata=target_metadata.get(name)
                     )
-            context.execute("--running migrations for engine %s" % name)
-            context.run_migrations(engine=name)
+            context.execute("-- running migrations for database %s" % name)
+            context.run_migrations(engine_name=name)
 
         if USE_TWOPHASE:
             for rec in engines.values():
diff --git a/alembic/templates/multidb/script.py.mako b/alembic/templates/multidb/script.py.mako
index 68d5e72..582492b 100644
--- a/alembic/templates/multidb/script.py.mako
+++ b/alembic/templates/multidb/script.py.mako
@@ -1,4 +1,7 @@
-"""${message}
+<%!
+import re
+
+%>"""${message}
 
 Revision ID: ${up_revision}
 Revises: ${down_revision}
@@ -14,21 +17,24 @@ from alembic import op
 import sqlalchemy as sa
 ${imports if imports else ""}
 
-def upgrade(engine):
-    eval("upgrade_%s" % engine.name)()
+def upgrade(engine_name):
+    eval("upgrade_%s" % engine_name)()
 
 
-def downgrade(engine):
-    eval("upgrade_%s" % engine.name)()
+def downgrade(engine_name):
+    eval("downgrade_%s" % engine_name)()
 
+<%
+    db_names = context.get("config").get_main_option("databases")
+%>
 
-% for engine in ["engine1", "engine2"]:
+% for db_name in re.split(r',\s*', db_names):
 
-def upgrade_${engine}():
-    ${context.get("%s_upgrades" % engine, "pass")}
+def upgrade_${db_name}():
+    ${context.get("%s_upgrades" % db_name, "pass")}
 
 
-def downgrade_${engine}():
-    ${context.get("%s_downgrades" % engine, "pass")}
+def downgrade_${db_name}():
+    ${context.get("%s_downgrades" % db_name, "pass")}
 
-% endfor
\ No newline at end of file
+% endfor