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: http://test.kernel.org/svn/autotest/trunk@4486 592f7852-d20e-0410-864c-8624ca9c26a4
diff --git a/database/db_utils.py b/database/db_utils.py
index 18f5495..7233265 100644
--- a/database/db_utils.py
+++ b/database/db_utils.py
@@ -2,6 +2,10 @@
 VIEW_TYPE = object()
 
 
+class NameMissingException(Exception):
+    pass
+
+
 def drop_views(manager, views):
     """
     Drops the specified views from the database
@@ -12,7 +16,7 @@
     @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 @@
     @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 @@
     @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 @@
     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 @@
 
     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 6c117da..72b7180 100755
--- a/database/db_utils_unittest.py
+++ b/database/db_utils_unittest.py
@@ -28,7 +28,7 @@
     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 b2fac89..a5d2ccc 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 @@
   '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 0000000..70c0eed
--- /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)