Fix migration 59's handling of a fresh database. The migration did not handle

the case where the Django auth tables did not exist. The permissions will be
inserted into the table later upon the first run of .
Also add a migration to properly create the permissions entries in the existing
tables.

Signed-off-by: James Ren <jamesren@google.com>



git-svn-id: svn://test.kernel.org/autotest/trunk@4486 592f7852-d20e-0410-864c-8624ca9c26a4

4 files changed, 75 insertions, 6 deletions

diff --git a/database/db_utils.py b/database/db_utils.py
index 18f5495f..72332659 100644
--- a/database/db_utils.py
+++ b/database/db_utils.py
@@ -2,6 +2,10 @@ TABLE_TYPE = object()
 VIEW_TYPE = object()
 
 
+class NameMissingException(Exception):
+    pass
+
+
 def drop_views(manager, views):
     """
     Drops the specified views from the database
@@ -12,7 +16,7 @@ def drop_views(manager, views):
     @param manager the migration manager
     @param views the views to drop
     """
-    _check_exists(manager, views, VIEW_TYPE)
+    check_exists(manager, views, VIEW_TYPE)
     for view in views:
         manager.execute('DROP VIEW `%s`' % view)
 
@@ -28,7 +32,7 @@ def rename(manager, mapping):
     @param mapping a dictionary of orig_name => new_name. Any table not matching
                    an entry in this dictionary will not be renamed
     """
-    _check_exists(manager, (table for table, _ in mapping.iteritems()),
+    check_exists(manager, (table for table, _ in mapping.iteritems()),
                   TABLE_TYPE)
     for orig_name, new_name in mapping.iteritems():
         manager.execute('RENAME TABLE `%s` TO `%s`' % (orig_name, new_name))
@@ -45,7 +49,7 @@ def move_tables(manager, src_manager, tables):
     @param src_manager a migration manager that handles the source database
     @param tables a list of tables to move
     """
-    _check_exists(src_manager, tables, TABLE_TYPE)
+    check_exists(src_manager, tables, TABLE_TYPE)
     for table in tables:
         manager.execute('RENAME TABLE `%(db)s`.`%(table)s` TO `%(table)s`'
                         % dict(db=src_manager.get_db_name(), table=table))
@@ -60,7 +64,7 @@ def drop_database(manager):
     manager.execute('DROP DATABASE `%s`' % manager.get_db_name())
 
 
-def _check_exists(manager, names, type):
+def check_exists(manager, names, type):
     """
     Checks if the tables or views exists.
 
@@ -84,4 +88,15 @@ def _check_exists(manager, names, type):
 
     for name in names:
         if name not in existing_names:
-            raise Exception('%s missing from database, stopping' % name)
+            raise NameMissingException(
+                    '%s missing from database, stopping' % name)
+
+
+DJANGO_AUTH_TABLES = ('auth_group', 'auth_group_permissions', 'auth_permission')
+
+def auth_tables_exist(manager):
+    try:
+        check_exists(manager, DJANGO_AUTH_TABLES, TABLE_TYPE)
+        return True
+    except NameMissingException:
+        return False
diff --git a/database/db_utils_unittest.py b/database/db_utils_unittest.py
index 6c117dab..72b7180e 100755
--- a/database/db_utils_unittest.py
+++ b/database/db_utils_unittest.py
@@ -28,7 +28,7 @@ class UtilsTest(unittest.TestCase):
     def test_check_exists(self):
         views = ('view1', 'view2')
         def _call_check_exists():
-            db_utils._check_exists(self.manager, views, db_utils.VIEW_TYPE)
+            db_utils.check_exists(self.manager, views, db_utils.VIEW_TYPE)
 
         self._setup_exists_expects(views, 'VIEWS')
         _call_check_exists()
diff --git a/frontend/migrations/059_drone_sets_permissions.py b/frontend/migrations/059_drone_sets_permissions.py
index b2fac896..a5d2cccd 100644
--- a/frontend/migrations/059_drone_sets_permissions.py
+++ b/frontend/migrations/059_drone_sets_permissions.py
@@ -1,3 +1,6 @@
+import common
+from autotest_lib.database import db_utils
+
 UP_SQL = """
 SET @group_id = (SELECT id FROM auth_group WHERE name = 'Basic Admin');
 
@@ -17,3 +20,20 @@ WHERE auth_group.name = 'Basic Admin' AND codename IN (
   'add_droneset', 'change_droneset', 'delete_droneset', 'add_drone',
   'change_drone', 'delete_drone');
 """
+
+
+def migrate_up(manager):
+    """
+    If the auth tables don't exist, we shouldn't try to set the permissions.
+
+    The auth tables will exist if this is an existing Autotest installation. If
+    they don't, then this is a fresh installation, and the user will run
+    `manage.py syncdb` later, which will add the proper permissions.
+    """
+    if db_utils.auth_tables_exist(manager):
+        manager.execute_script(UP_SQL)
+
+
+def migrate_down(manager):
+    if db_utils.auth_tables_exist(manager):
+        manager.execute_script(DOWN_SQL)
diff --git a/frontend/migrations/061_drone_sets_permissions_proper.py b/frontend/migrations/061_drone_sets_permissions_proper.py
new file mode 100644
index 00000000..70c0eede
--- /dev/null
+++ b/frontend/migrations/061_drone_sets_permissions_proper.py
@@ -0,0 +1,34 @@
+from django.core import management
+import common
+from autotest_lib.frontend import settings
+from autotest_lib.database import db_utils
+
+AFE_MIGRATION_NAME = '059_drone_sets_permissions'
+migrations_module = __import__('autotest_lib.frontend.migrations', globals(),
+                               locals(), [AFE_MIGRATION_NAME])
+migration_059 = getattr(migrations_module, AFE_MIGRATION_NAME)
+
+
+def migrate_up(manager):
+    """
+    If the auth tables don't exist, we shouldn't try to set the permissions.
+
+    See migration 059
+    """
+    if db_utils.auth_tables_exist(manager):
+        management.setup_environ(settings)
+        # These have to be imported after the environment is set up
+        from django.contrib.contenttypes import management as content_management
+        from django.contrib.auth import management as auth_management
+        from django.db import models as db_models
+
+        content_management.update_all_contenttypes()
+        for app in db_models.get_apps():
+            auth_management.create_permissions(app, None, 2)
+
+        manager.execute_script(migration_059.UP_SQL)
+
+
+def migrate_down(manager):
+    if db_utils.auth_tables_exist(manager):
+        manager.execute_script(migration_059.DOWN_SQL)