From 71f51f67ac9de9b1897a7ae889f6f6a0cf1c0c81 Mon Sep 17 00:00:00 2001
From: Bryce Lohr <bryce.lohr@gmail.com>
Date: Fri, 27 Apr 2012 00:47:14 -0400
Subject: [PATCH] Updated the 'revision' command to pass the current Alembic
 config to the script templates, so it could use it to generate blocks based
 on current config. Useful for multidb configurations, where you want each
 database accounted for in the migrations.

---
 alembic/command.py                       |  4 +++-
 alembic/templates/multidb/script.py.mako | 18 ++++++++++++------
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/alembic/command.py b/alembic/command.py
index ed7b8302..67b0d668 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/templates/multidb/script.py.mako b/alembic/templates/multidb/script.py.mako
index 7aef96f7..582492bb 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}
@@ -21,14 +24,17 @@ def upgrade(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
-- 
2.47.2