From 08cd10d84fd89df3f0cad3835ba8ab8b8000d4b2 Mon Sep 17 00:00:00 2001
From: tilly-Q <nattilypigeonfowl@gmail.com>
Date: Mon, 29 Jul 2013 17:15:29 -0400
Subject: [PATCH] I actually had to do a bit more work than I thought, because
 I needed to account for plugins. In this commit I changed the
 MigrationManager and DatabaseData ob- jects to account for FOUNDATIONS in any
 plugin's (or main program's) models.py file.

---
 mediagoblin/db/migration_tools.py    | 15 +++++++--------
 mediagoblin/gmg_commands/dbupdate.py | 25 ++++++++++++++++++++-----
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/mediagoblin/db/migration_tools.py b/mediagoblin/db/migration_tools.py
index 1192836d..ad137683 100644
--- a/mediagoblin/db/migration_tools.py
+++ b/mediagoblin/db/migration_tools.py
@@ -29,7 +29,7 @@ class MigrationManager(object):
     to the latest migrations, etc.
     """
 
-    def __init__(self, name, models, migration_registry, session,
+    def __init__(self, name, models, foundations, migration_registry, session,
                  printer=simple_printer):
         """
         Args:
@@ -40,6 +40,7 @@ class MigrationManager(object):
         """
         self.name = unicode(name)
         self.models = models
+        self.foundations = foundations
         self.session = session
         self.migration_registry = migration_registry
         self._sorted_migrations = None
@@ -145,12 +146,11 @@ class MigrationManager(object):
         Create the table foundations (default rows) as layed out in FOUNDATIONS
             in mediagoblin.db.models
         """
-        from mediagoblin.db.models import FOUNDATIONS as MAIN_FOUNDATIONS
-        for Model, rows in MAIN_FOUNDATIONS.items():
-            print u'\n--> Laying foundations for %s table' % Model.__name__
+        for Model, rows in self.foundations.items():
+            print u'\n  + Laying foundations for %s table' % (Model.__name__)
             for parameters in rows:
-                row = Model(**parameters)
-                row.save()
+                new_row = Model(**parameters)
+                new_row.save()
 
     def create_new_migration_record(self):
         """
@@ -215,8 +215,7 @@ class MigrationManager(object):
             self.init_tables()
             # auto-set at latest migration number
             self.create_new_migration_record()
-            if self.name==u'__main__':
-                self.populate_table_foundations()
+            self.populate_table_foundations()
             
             self.printer(u"done.\n")
             self.set_current_migration()
diff --git a/mediagoblin/gmg_commands/dbupdate.py b/mediagoblin/gmg_commands/dbupdate.py
index 00007567..bad3e352 100644
--- a/mediagoblin/gmg_commands/dbupdate.py
+++ b/mediagoblin/gmg_commands/dbupdate.py
@@ -32,14 +32,15 @@ def dbupdate_parse_setup(subparser):
 
 
 class DatabaseData(object):
-    def __init__(self, name, models, migrations):
+    def __init__(self, name, models, foundations, migrations):
         self.name = name
         self.models = models
+        self.foundations = foundations
         self.migrations = migrations
 
     def make_migration_manager(self, session):
         return MigrationManager(
-            self.name, self.models, self.migrations, session)
+            self.name, self.models, self.foundations, self.migrations, session)
 
 
 def gather_database_data(plugins):
@@ -54,10 +55,11 @@ def gather_database_data(plugins):
     # Add main first
     from mediagoblin.db.models import MODELS as MAIN_MODELS
     from mediagoblin.db.migrations import MIGRATIONS as MAIN_MIGRATIONS
+    from mediagoblin.db.models import FOUNDATIONS as MAIN_FOUNDATIONS
 
     managed_dbdata.append(
         DatabaseData(
-            u'__main__', MAIN_MODELS, MAIN_MIGRATIONS))
+            u'__main__', MAIN_MODELS, MAIN_FOUNDATIONS, MAIN_MIGRATIONS))
 
     for plugin in plugins:
         try:
@@ -83,13 +85,26 @@ forgotten to add it? ({1})'.format(plugin, exc))
 
             migrations = {}
         except AttributeError as exc:
-            _log.debug('Cloud not find MIGRATIONS in {0}.migrations, have you \
+            _log.debug('Could not find MIGRATIONS in {0}.migrations, have you \
 forgotten to add it? ({1})'.format(plugin, exc))
             migrations = {}
 
+        try:
+            foundations = import_component('{0}.models:FOUNDATIONS'.format(plugin))
+        except ImportError as exc:
+            _log.debug('No foundations found for {0}: {1}'.format(
+                plugin,
+                exc))
+
+            foundations = []
+        except AttributeError as exc:
+            _log.debug('Could not find FOUNDATIONS in {0}.models, have you \
+forgotten to add it? ({1})'.format(plugin, exc))
+            foundations = {}
+
         if models:
             managed_dbdata.append(
-                    DatabaseData(plugin, models, migrations))
+                    DatabaseData(plugin, models, foundations, migrations))
 
 
     return managed_dbdata
-- 
2.25.1