[autotest] Use global database for tko models in django v3.

This adds the global database to django.
A django database router is added to determine to which database
should be used for which models. All tko models are always taken or
written from or into the global database while all other objects
remain unchanged.

BUG=chromium:422637
TEST=Ran suites, syncdb, restart apache, ran a suite.
Confirm global database connection is closed after each attempt to get
test view: https://x20web.corp.google.com/~beeps/log/db_routers.html

Change-Id: Idf6933d1d112bbc5a2896fa61afd03f6604dafb5
Reviewed-on: https://chromium-review.googlesource.com/223501
Tested-by: Dan Shi <[email protected]>
Reviewed-by: Prashanth B <[email protected]>
Commit-Queue: Dan Shi <[email protected]>
diff --git a/frontend/db_router.py b/frontend/db_router.py
new file mode 100644
index 0000000..649fb90
--- /dev/null
+++ b/frontend/db_router.py
@@ -0,0 +1,92 @@
+# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Django database Router
+
+Django gets configured with three database connections in frontend/settings.py.
+- The default database
+    - This database should be used for most things.
+    - For the master, this is the global database.
+    - For shards, this this is the shard-local database.
+- The global database
+    - For the master, this is the same database as default, which is the global
+      database.
+    - For the shards, this is the global database (the same as for the master).
+- The readonly connection
+    - This should be the same database as the global database, but it should
+      use an account on the database that only has readonly permissions.
+
+The reason shards need two distinct databases for different objects is, that
+the tko parser should always write to the global database. Otherwise test
+results wouldn't be synced back to the master and would not be accessible in one
+place.
+
+Therefore this class will route all queries that involve `tko_`-prefixed tables
+to the global database. For all others this router will not give a hint, which
+means the default database will be used.
+"""
+
+class Router(object):
+    """
+    Decide if an object should be written to the default or to the global db.
+
+    This is an implementaton of Django's multi-database router interface:
+    https://docs.djangoproject.com/en/1.5/topics/db/multi-db/
+    """
+
+    def _should_be_in_global(self, model):
+        """Returns True if the model should be stored in the global db.
+
+        @param model: Model to decide for.
+
+        @return: True if querying the model requires global database.
+        """
+        return model._meta.db_table.startswith('tko_')
+
+
+    def db_for_read(self, model, **hints):
+        """
+        Decides if the global database should be used for a reading access.
+
+        @param model: Model to decide for.
+        @param hints: Optional arguments to determine which database for read.
+
+        @returns: 'global' for all tko models, None otherwise. None means the
+                  router doesn't have an opinion.
+        """
+        if self._should_be_in_global(model):
+            return 'global'
+        return None
+
+
+    def db_for_write(self, model, **hints):
+        """
+        Decides if the global database should be used for a writing access.
+
+        @param model: Model to decide for.
+        @param hints: Optional arguments to determine which database for write.
+
+        @returns: 'global' for all tko models, None otherwise. None means the
+                  router doesn't have an opinion.
+        """
+        if self._should_be_in_global(model):
+            return 'global'
+        return None
+
+
+    def allow_relation(self, obj1, obj2, **hints):
+        """
+        Allow relations only if either both are in tko_ tables or none is.
+
+        @param obj1: First object involved in the relation.
+        @param obj2: Second object involved in the relation.
+        @param hints: Optional arguments to determine if relation is allowed.
+
+        @returns False, if the relation should be prohibited,
+                 None, if the router doesn't have an opinion.
+        """
+        if (not self._should_be_in_global(type(obj1)) ==
+            self._should_be_in_global(type(obj2))):
+            return False
+        return None
diff --git a/frontend/settings.py b/frontend/settings.py
index 95ca552..5dad550 100644
--- a/frontend/settings.py
+++ b/frontend/settings.py
@@ -1,4 +1,14 @@
 """Django settings for frontend project.
+
+Two databases are configured for the use with django here. One for tko tables,
+which will always be the same database for all instances (the global database),
+and one for everything else, which will be the same as the global database for
+the master, but a local database for shards.
+Additionally there is a third database connection for read only access to the
+global database.
+
+This is implemented using a Django database router.
+For more details on how the routing works, see db_router.py.
 """
 
 import os
@@ -28,6 +38,7 @@
 ALLOWED_HOSTS = '*'
 
 DATABASES = {'default': AUTOTEST_DEFAULT,
+             'global': AUTOTEST_GLOBAL,
              'readonly': AUTOTEST_READONLY,}
 
 # Have to set SECRET_KEY before importing connections because of this bug:
@@ -36,6 +47,11 @@
 # Make this unique, and don't share it with anybody.
 SECRET_KEY = 'pn-t15u(epetamdflb%dqaaxw+5u&2#0u-jah70w1l*_9*)=n7'
 
+# Do not do this here or from the router, or most unit tests will fail.
+# from django.db import connection
+
+DATABASE_ROUTERS = ['autotest_lib.frontend.db_router.Router']
+
 # prefix applied to all URLs - useful if requests are coming through apache,
 # and you need this app to coexist with others
 URL_PREFIX = 'afe/server/'
diff --git a/frontend/setup_django_environment.py b/frontend/setup_django_environment.py
index 055e9e9..cd390a3 100644
--- a/frontend/setup_django_environment.py
+++ b/frontend/setup_django_environment.py
@@ -4,7 +4,20 @@
 
 os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'autotest_lib.frontend.settings')
 
+def _enable_autocommit_by_name(name):
+    """Enable autocommit for the connection with matching name.
+
+    @param name: Name of the connection.
+    """
+    from django.db import connections
+    # ensure a connection is open
+    connections[name].cursor()
+    connections[name].connection.autocommit(True)
+
+
 def enable_autocommit():
-    from django.db import connection
-    connection.cursor() # ensure a connection is open
-    connection.connection.autocommit(True)
+    """Enable autocommit for default and global connection.
+    """
+    _enable_autocommit_by_name('default')
+    _enable_autocommit_by_name('global')
+
diff --git a/frontend/setup_test_environment.py b/frontend/setup_test_environment.py
index 9fe58f5..45b76f9 100644
--- a/frontend/setup_test_environment.py
+++ b/frontend/setup_test_environment.py
@@ -1,4 +1,3 @@
-import tempfile, shutil, os
 from django.core import management
 from django.conf import settings
 import common
@@ -12,6 +11,11 @@
     'autotest_lib.frontend.db.backends.afe_sqlite')
 settings.DATABASES['default']['NAME'] = ':memory:'
 
+settings.DATABASES['global'] = {}
+settings.DATABASES['global']['ENGINE'] = (
+    'autotest_lib.frontend.db.backends.afe_sqlite')
+settings.DATABASES['global']['NAME'] = ':memory:'
+
 settings.DATABASES['readonly'] = {}
 settings.DATABASES['readonly']['ENGINE'] = (
     'autotest_lib.frontend.db.backends.afe_sqlite')
@@ -22,21 +26,32 @@
 
 connection = connections['default']
 connection_readonly = connections['readonly']
+connection_global = connections['global']
 
 def run_syncdb(verbosity=0):
+    """Call syncdb command to make sure database schema is uptodate.
+
+    @param verbosity: Level of verbosity of the command, default to 0.
+    """
     management.call_command('syncdb', verbosity=verbosity, interactive=False)
     management.call_command('syncdb', verbosity=verbosity, interactive=False,
                              database='readonly')
+    management.call_command('syncdb', verbosity=verbosity, interactive=False,
+                             database='global')
+
 
 def destroy_test_database():
+    """Close all connection to the test database.
+    """
     connection.close()
     connection_readonly.close()
+    connection_global.close()
     # Django brilliantly ignores close() requests on in-memory DBs to keep us
     # naive users from accidentally destroying data.  So reach in and close
     # the real connection ourselves.
     # Note this depends on Django internals and will likely need to be changed
     # when we upgrade Django.
-    for con in [connection, connection_readonly]:
+    for con in [connection, connection_global, connection_readonly]:
         real_connection = con.connection
         if real_connection is not None:
             real_connection.close()
@@ -44,11 +59,15 @@
 
 
 def set_up():
+    """Run setup before test starts.
+    """
     run_syncdb()
     readonly_connection.set_globally_disabled(True)
 
 
 def tear_down():
+    """Run cleanup after test is completed.
+    """
     readonly_connection.set_globally_disabled(False)
     destroy_test_database()
 
diff --git a/tko/site_parse_unittest.py b/tko/site_parse_unittest.py
index 82925a6..9a98f65 100755
--- a/tko/site_parse_unittest.py
+++ b/tko/site_parse_unittest.py
@@ -7,15 +7,24 @@
 
 #pylint: disable-msg=C0111
 
-import os, shutil, tempfile, unittest
+import mox, os, shutil, tempfile, unittest
+
+from django.conf import settings
 
 import common
 from autotest_lib.client.common_lib import global_config
 from autotest_lib.frontend import database_settings_helper
 from autotest_lib.frontend import setup_django_environment
 from autotest_lib.frontend import setup_test_environment
+from autotest_lib.frontend.afe import frontend_test_utils
+from autotest_lib.frontend.afe import models as django_afe_models
+from autotest_lib.frontend.tko import models as django_tko_models
+from autotest_lib.tko import db as tko_db
 from autotest_lib.tko.site_parse import StackTrace
 
+# Have to import this after setup_django_environment and setup_test_environment.
+# It creates a database connection, so the mocking has to be done first.
+from django.db import connections
 
 class stack_trace_test(unittest.TestCase):
 
@@ -107,24 +116,111 @@
         self.assertEqual(version, '1166.0.0')
 
 
+class database_selection_test(mox.MoxTestBase,
+                              frontend_test_utils.FrontendTestMixin):
+
+    def setUp(self):
+        super(database_selection_test, self).setUp()
+        self._frontend_common_setup(fill_data=False)
+
+
+    def tearDown(self):
+        super(database_selection_test, self).tearDown()
+        self._frontend_common_teardown()
+        global_config.global_config.reset_config_values()
+
+
+    def assertQueries(self, database, assert_in, assert_not_in):
+        assert_in_found = False
+        for query in connections[database].queries:
+            sql = query['sql']
+            # Ignore CREATE TABLE statements as they are always executed
+            if 'INSERT INTO' in sql or 'SELECT' in sql:
+                self.assertNotIn(assert_not_in, sql)
+                if assert_in in sql:
+                    assert_in_found = True
+        self.assertTrue(assert_in_found)
+
+
+    def testDjangoModels(self):
+        # If DEBUG=False connection.query will be empty
+        settings.DEBUG = True
+
+        afe_job = django_afe_models.Job.objects.create(created_on='2014-08-12')
+        # Machine has less dependencies than tko Job so it's easier to create
+        tko_job = django_tko_models.Machine.objects.create()
+
+        django_afe_models.Job.objects.get(pk=afe_job.id)
+        django_tko_models.Machine.objects.get(pk=tko_job.pk)
+
+        self.assertQueries('global', 'tko_machines', 'afe_jobs')
+        self.assertQueries('default', 'afe_jobs', 'tko_machines')
+
+        # Avoid unnecessary debug output from other tests
+        settings.DEBUG = True
+
+
     def testRunOnShardWithoutGlobalConfigsFails(self):
         global_config.global_config.override_config_value(
                 'SHARD', 'shard_hostname', 'host1')
+        from autotest_lib.frontend import settings
         # settings module was already loaded during the imports of this file,
         # so before the configuration setting was made, therefore reload it:
         reload(database_settings_helper)
         self.assertRaises(global_config.ConfigError,
-                          database_settings_helper.get_global_db_config)
+                          reload, settings)
 
 
     def testRunOnMasterWithoutGlobalConfigsWorks(self):
         global_config.global_config.override_config_value(
                 'SHARD', 'shard_hostname', '')
-        from autotest_lib.frontend import database_settings_helper
+        from autotest_lib.frontend import settings
         # settings module was already loaded during the imports of this file,
         # so before the configuration setting was made, therefore reload it:
         reload(database_settings_helper)
-        database_settings_helper.get_global_db_config()
+        reload(settings)
+
+
+    def testTkoDatabase(self):
+        global_host = 'GLOBAL_HOST'
+        global_user = 'GLOBAL_USER'
+        global_db = 'GLOBAL_DB'
+        global_pw = 'GLOBAL_PW'
+        global_port = ''
+        local_host = 'LOCAL_HOST'
+
+        global_config.global_config.override_config_value(
+                'AUTOTEST_WEB', 'global_db_type', '')
+
+        global_config.global_config.override_config_value(
+                'AUTOTEST_WEB', 'global_db_host', global_host)
+        global_config.global_config.override_config_value(
+                'AUTOTEST_WEB', 'global_db_database', global_db)
+        global_config.global_config.override_config_value(
+                'AUTOTEST_WEB', 'global_db_user', global_user)
+        global_config.global_config.override_config_value(
+                'AUTOTEST_WEB', 'global_db_password', global_pw)
+        global_config.global_config.override_config_value(
+                'AUTOTEST_WEB', 'host', local_host)
+
+        class ConnectCalledException(Exception):
+            pass
+
+        # We're only interested in the parameters connect is called with here.
+        # Take the fast path out so we don't have to mock all the other calls
+        # that will later be made on the connection
+        def fake_connect(*args, **kwargs):
+            raise ConnectCalledException
+
+        tko_db.db_sql.connect = None
+        self.mox.StubOutWithMock(tko_db.db_sql, 'connect')
+        tko_db.db_sql.connect(
+                global_host, global_db, global_user, global_pw,
+                global_port).WithSideEffects(fake_connect)
+
+        self.mox.ReplayAll()
+
+        self.assertRaises(ConnectCalledException, tko_db.db_sql)
 
 
 if __name__ == "__main__":
diff --git a/utils/unittest_suite.py b/utils/unittest_suite.py
index 9b8c121..277a9a9 100755
--- a/utils/unittest_suite.py
+++ b/utils/unittest_suite.py
@@ -46,6 +46,7 @@
         'rdb_cache_unittests.py',
         'scheduler_lib_unittest.py',
         'host_scheduler_unittests.py',
+        'site_parse_unittest.py',
         ))
 
 REQUIRES_MYSQLDB = set((