1 files changed, 71 insertions, 0 deletions
diff --git a/mediagoblin/gmg_commands/alembic_commands.py b/mediagoblin/gmg_commands/alembic_commands.py
new file mode 100644
index 00000000..4e73ec13
--- /dev/null
+++ b/mediagoblin/gmg_commands/alembic_commands.py
@@ -0,0 +1,71 @@
+# GNU MediaGoblin -- federated, autonomous media hosting
+# Copyright (C) 2011, 2012 MediaGoblin contributors.  See AUTHORS.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import argparse
+
+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
+from mediagoblin.tools.common import import_component
+
+
+class FudgedCommandLine(config.CommandLine):
+    def main(self, args, db, global_config):
+        options = self.parser.parse_args(args.args_for_alembic)
+
+        # This is useful for people running commands involving plugins, such
+        # as building autogenerated revisions.
+        if args.with_plugins:
+            plugins = global_config.get('plugins', {}).keys()
+            for plugin in plugins:
+                try:
+                    import_component('{0}.models:MODELS'.format(plugin))
+                except ImportError:
+                    # It doesn't really matter if there's no models to import
+                    # here.
+                    pass
+
+        # This code is inspired by a hack in Alembic, but isn't the same really.
+        # Regardless, Alembic is Expat licensed.
+        if not hasattr(options, "cmd"):
+            print(
+                "* Only use this command if you know what you are doing! *\n"
+                "If not, use the 'gmg dbupdate' command instead.\n\n"
+                "(You can also pass --with-plugins here.)"
+                "Alembic help:\n")
+            self.parser.print_help()
+            return
+        else:
+            Session = sessionmaker(bind=db.engine)
+            session = Session()
+
+            cfg = build_alembic_config(global_config, options, session)
+
+            self.run_cmd(cfg, options)
+        
+def parser_setup(subparser):
+    subparser.add_argument(
+        "--with-plugins", action="store_true",
+        help="Import all plugins' models before running alembic commands.")
+    subparser.add_argument("args_for_alembic", nargs=argparse.REMAINDER)
+
+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, global_config)