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.

1 files changed, 12 insertions, 6 deletions

diff --git a/alembic/templates/multidb/script.py.mako b/alembic/templates/multidb/script.py.mako
index 7aef96f..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}
@@ -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