Add build_alembic_config, use it to add plugin migrations to alembic config

2 files changed, 40 insertions, 9 deletions

diff --git a/mediagoblin/db/migration_tools.py b/mediagoblin/db/migration_tools.py
index bd242894..c6beba8b 100644
--- a/mediagoblin/db/migration_tools.py
+++ b/mediagoblin/db/migration_tools.py
@@ -18,6 +18,7 @@ from __future__ import unicode_literals
 
 import logging
 import os
+import pkg_resources
 
 from alembic import command
 from alembic.config import Config
@@ -401,3 +402,36 @@ def model_iteration_hack(db, query):
     return db.execute(query)
 
 
+def build_alembic_config(global_config, cmd_options, session):
+    """
+    Build up a config that the alembic tooling can use based on our
+    configuration.  Initialize the database session appropriately
+    as well.
+    """
+    root_dir = os.path.abspath(os.path.dirname(os.path.dirname(
+        os.path.dirname(__file__))))
+    alembic_cfg_path = os.path.join(root_dir, 'alembic.ini')
+    cfg = Config(alembic_cfg_path,
+                 cmd_opts=cmd_options)
+    cfg.attributes["session"] = session
+
+    version_locations = [
+        pkg_resources.resource_filename(
+            "mediagoblin.db", os.path.join("migrations", "versions")),
+    ]
+
+    cfg.set_main_option("sqlalchemy.url", str(session.get_bind().url))
+
+    for plugin in global_config.get("plugins", []):
+        plugin_migrations = pkg_resources.resource_filename(
+            plugin, "migrations")
+        is_migrations_dir = (os.path.exists(plugin_migrations) and
+                             os.path.isdir(plugin_migrations))
+        if is_migrations_dir:
+            version_locations.append(plugin_migrations)
+
+    cfg.set_main_option(
+        "version_locations",
+        " ".join(version_locations))
+
+    return cfg
diff --git a/mediagoblin/gmg_commands/alembic_commands.py b/mediagoblin/gmg_commands/alembic_commands.py
index f255af73..58b931dc 100644
--- a/mediagoblin/gmg_commands/alembic_commands.py
+++ b/mediagoblin/gmg_commands/alembic_commands.py
@@ -15,17 +15,17 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import argparse
-import os
 
 from alembic import config
 from sqlalchemy.orm import sessionmaker
 
 from mediagoblin.db.open import setup_connection_and_db_from_config
 from mediagoblin.init import setup_global_and_app_config
+from mediagoblin.db.migration_tools import build_alembic_config
 
 
 class FudgedCommandLine(config.CommandLine):
-    def main(self, args, db):
+    def main(self, args, db, global_config):
         options = self.parser.parse_args(args.args_for_alembic)
         # This code is inspired by a hack in Alembic, but isn't the same really.
         # Regardless, Alembic is Expat licensed.
@@ -38,13 +38,10 @@ class FudgedCommandLine(config.CommandLine):
             return
         else:
             Session = sessionmaker(bind=db.engine)
+            session = Session()
+
+            cfg = build_alembic_config(global_config, options, session)
 
-            root_dir = os.path.abspath(os.path.dirname(os.path.dirname(
-                os.path.dirname(__file__))))
-            alembic_cfg_path = os.path.join(root_dir, 'alembic.ini')
-            cfg = config.Config(alembic_cfg_path,
-                                cmd_opts=options)
-            cfg.attributes["session"] = Session()
             self.run_cmd(cfg, options)
         
 def parser_setup(subparser):
@@ -53,4 +50,4 @@ def parser_setup(subparser):
 def raw_alembic_cli(args):
     global_config, app_config = setup_global_and_app_config(args.conf_file)
     db = setup_connection_and_db_from_config(app_config, migrations=False)
-    FudgedCommandLine().main(args, db)
+    FudgedCommandLine().main(args, db, global_config)