[autotest] Refactor readonly_connection
The way readonly connections were used used
some monkey patching and manual connection handling.
Code gets easier to maintain if Django manages all the
database connections and no monkeypatching is done.
This is what this CL achieves.
BUG=chromium:422637
DEPLOY=apache,scheduler,host_scheduler,shard_client
TEST=Ran suites.
Change-Id: I2761af10f5172bb3558e7dbd7c1f1875f7a7cb51
Reviewed-on: https://chromium-review.googlesource.com/223393
Tested-by: Jakob Jülich <jakobjuelich@chromium.org>
Reviewed-by: Dan Shi <dshi@chromium.org>
Reviewed-by: Fang Deng <fdeng@chromium.org>
Commit-Queue: Jakob Jülich <jakobjuelich@chromium.org>
diff --git a/frontend/afe/model_logic.py b/frontend/afe/model_logic.py
index 0c2bafc..2566e24 100644
--- a/frontend/afe/model_logic.py
+++ b/frontend/afe/model_logic.py
@@ -5,12 +5,11 @@
import re
import time
import django.core.exceptions
-from django.db import models as dbmodels, backend, connection
+from django.db import models as dbmodels, backend, connection, connections
from django.db.models.sql import query
import django.db.models.sql.where
from django.utils import datastructures
from autotest_lib.frontend.afe import rdb_model_extensions
-from autotest_lib.frontend.afe import readonly_connection
class ValidationError(django.core.exceptions.ValidationError):
@@ -19,79 +18,11 @@
value is a dictionary mapping field names to error strings.
"""
-
-def _wrap_with_readonly(method):
- def wrapper_method(*args, **kwargs):
- readonly_connection.connection().set_django_connection()
- try:
- return method(*args, **kwargs)
- finally:
- readonly_connection.connection().unset_django_connection()
- wrapper_method.__name__ = method.__name__
- return wrapper_method
-
-
def _quote_name(name):
"""Shorthand for connection.ops.quote_name()."""
return connection.ops.quote_name(name)
-def _wrap_generator_with_readonly(generator):
- """
- We have to wrap generators specially. Assume it performs
- the query on the first call to next().
- """
- def wrapper_generator(*args, **kwargs):
- generator_obj = generator(*args, **kwargs)
- readonly_connection.connection().set_django_connection()
- try:
- first_value = generator_obj.next()
- finally:
- readonly_connection.connection().unset_django_connection()
- yield first_value
-
- while True:
- yield generator_obj.next()
-
- wrapper_generator.__name__ = generator.__name__
- return wrapper_generator
-
-
-def _make_queryset_readonly(queryset):
- """
- Wrap all methods that do database queries with a readonly connection.
- """
- db_query_methods = ['count', 'get', 'get_or_create', 'latest', 'in_bulk',
- 'delete']
- for method_name in db_query_methods:
- method = getattr(queryset, method_name)
- wrapped_method = _wrap_with_readonly(method)
- setattr(queryset, method_name, wrapped_method)
-
- queryset.iterator = _wrap_generator_with_readonly(queryset.iterator)
-
-
-class ReadonlyQuerySet(dbmodels.query.QuerySet):
- """
- QuerySet object that performs all database queries with the read-only
- connection.
- """
- def __init__(self, model=None, *args, **kwargs):
- super(ReadonlyQuerySet, self).__init__(model, *args, **kwargs)
- _make_queryset_readonly(self)
-
-
- def values(self, *fields):
- return self._clone(klass=ReadonlyValuesQuerySet,
- setup=True, _fields=fields)
-
-
-class ReadonlyValuesQuerySet(dbmodels.query.ValuesQuerySet):
- def __init__(self, model=None, *args, **kwargs):
- super(ReadonlyValuesQuerySet, self).__init__(model, *args, **kwargs)
- _make_queryset_readonly(self)
-
-
class LeasedHostManager(dbmodels.Manager):
"""Query manager for unleased, unlocked hosts.
"""
@@ -362,6 +293,13 @@
def _custom_select_query(self, query_set, selects):
+ """Execute a custom select query.
+
+ @param query_set: query set as returned by query_objects.
+ @param selects: Tables/Columns to select, e.g. tko_test_labels_list.id.
+
+ @returns: Result of the query as returned by cursor.fetchall().
+ """
compiler = query_set.query.get_compiler(using=query_set.db)
sql, params = compiler.as_sql()
from_ = sql[sql.find(' FROM'):]
@@ -372,7 +310,8 @@
distinct = ''
sql_query = ('SELECT ' + distinct + ','.join(selects) + from_)
- cursor = readonly_connection.connection().cursor()
+ # Chose the connection that's responsible for this type of object
+ cursor = connections[query_set.db].cursor()
cursor.execute(sql_query, params)
return cursor.fetchall()
@@ -458,7 +397,7 @@
def _query_pivot_table(self, base_objects_by_id, pivot_table,
- pivot_from_field, pivot_to_field):
+ pivot_from_field, pivot_to_field, related_model):
"""
@param id_list list of IDs of self.model objects to include
@param pivot_table the name of the pivot table
@@ -466,6 +405,8 @@
self.model
@param pivot_to_field a field name on pivot_table referencing the
related model.
+ @param related_model the related model
+
@returns pivot list of IDs (base_id, related_id)
"""
query = """
@@ -477,7 +418,12 @@
table=pivot_table,
id_list=','.join(str(id_) for id_
in base_objects_by_id.iterkeys()))
- cursor = readonly_connection.connection().cursor()
+
+ # Chose the connection that's responsible for this type of object
+ # The databases for related_model and the current model will always
+ # be the same, related_model is just easier to obtain here because
+ # self is only a ExtendedManager, not the object.
+ cursor = connections[related_model.objects.db].cursor()
cursor.execute(query)
return cursor.fetchall()
@@ -491,7 +437,8 @@
@returns a pivot iterator - see _get_pivot_iterator()
"""
id_pivot = self._query_pivot_table(base_objects_by_id, pivot_table,
- pivot_from_field, pivot_to_field)
+ pivot_from_field, pivot_to_field,
+ related_model)
all_related_ids = list(set(related_id for base_id, related_id
in id_pivot))
@@ -826,7 +773,9 @@
extra_args.setdefault('where', []).append(extra_where)
if extra_args:
query = query.extra(**extra_args)
- query = query._clone(klass=ReadonlyQuerySet)
+ # TODO: Use readonly connection for these queries.
+ # This has been disabled, because it's not used anyway, as the
+ # configured readonly user is the same as the real user anyway.
if apply_presentation:
query = cls.apply_presentation(query, filter_data)
diff --git a/frontend/afe/readonly_connection.py b/frontend/afe/readonly_connection.py
index 38f8c7d..b8dad9c 100644
--- a/frontend/afe/readonly_connection.py
+++ b/frontend/afe/readonly_connection.py
@@ -1,136 +1,36 @@
+# 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.
+
+"""This files contains helper methods to interact with the readonly database."""
+
+# Please note this file doesn't contain and should not contain any logic to
+# establish connections outside of Django, as that might lead to effects where
+# connections get leaked, which will lead to Django not cleaning them up
+# properly. See http://crbug.com/422637 for more details on this failure.
+
from django import db as django_db
-from django.conf import settings
-from django.core import signals
-class ReadOnlyConnection(object):
+_DISABLED = False
+
+def set_globally_disabled(disable):
+ """Disable and enable the use of readonly connections globally.
+
+ If disabled, connection() will return the global connection instead of the
+ readonly connection.
+
+ @param disable: When True, readonly connections will be disabled.
"""
- This class constructs a new connection to the DB using the read-only
- credentials from settings. It reaches into some internals of
- django.db.connection which are undocumented as far as I know, but I believe
- it works across many, if not all, of the backends.
- """
- _the_instance = None
-
- # support singleton
- @classmethod
- def get_connection(cls):
- if cls._the_instance is None:
- cls._the_instance = ReadOnlyConnection()
- return cls._the_instance
+ _DISABLED = disable
- @classmethod
- def set_globally_disabled(cls, disabled):
- """
- When globally disabled, the ReadOnlyConnection will simply pass through
- to the global Django connection.
- """
- if disabled:
- cls._the_instance = DummyReadOnlyConnection()
- else:
- cls._the_instance = None
-
-
- def __init__(self):
- self._connection = None
-
-
- def _open_connection(self):
- if self._connection is not None:
- return
- self._save_django_state()
- self._connection = self._get_readonly_connection()
- self._restore_django_state()
-
-
- def _save_django_state(self):
- self._old_connection = django_db.connection.connection
- _default_db = settings.DATABASES['default']
- self._old_host = _default_db['HOST']
- self._old_username = _default_db['USER']
- self._old_password = _default_db['PASSWORD']
-
-
- def _restore_django_state(self):
- django_db.connection.connection = self._old_connection
- _default_db = settings.DATABASES['default']
- _default_db['HOST'] = self._old_host
- _default_db['USER'] = self._old_username
- _default_db['PASSWORD'] = self._old_password
-
-
- def _get_readonly_connection(self):
- _default_db = settings.DATABASES['default']
- _default_db['HOST'] = _default_db['READONLY_HOST']
- _default_db['USER'] = _default_db['READONLY_USER']
- _default_db['PASSWORD'] = _default_db['READONLY_PASSWORD']
- reload(django_db)
- # cursor() causes a new connection to be created
- cursor = django_db.connection.cursor()
- assert django_db.connection.connection is not None
- return django_db.connection.connection
-
-
- def set_django_connection(self):
- assert (django_db.connection.connection != self._connection or
- self._connection is None)
- self._open_connection()
- self._old_connection = django_db.connection.connection
- django_db.connection.connection = self._connection
-
-
- def unset_django_connection(self):
- assert self._connection is not None
- assert django_db.connection.connection == self._connection
- django_db.connection.connection = self._old_connection
-
-
- def cursor(self):
- self._open_connection()
- return self._connection.cursor()
-
-
- def close(self):
- if self._connection is not None:
- assert django_db.connection.connection != self._connection
- self._connection.close()
- self._connection = None
-
-
-class DummyReadOnlyConnection(object):
- """
- A dummy version which passes queries straight to the global Django
- connection.
- """
-
- def __init__(self):
- self._is_set = False
-
-
- def set_django_connection(self):
- assert not self._is_set
- self._is_set = True
-
-
- def unset_django_connection(self):
- assert self._is_set
- self._is_set = False
-
-
- def cursor(self):
- return django_db.connection.cursor()
-
-
- def close(self):
- pass
-
-
-# convenience
def connection():
- return ReadOnlyConnection.get_connection()
+ """Return a readonly database connection."""
+ if _DISABLED:
+ return django_db.connections['global']
+ return django_db.connections['readonly']
-# close any open connection when request finishes
-def _close_connection(**unused_kwargs):
- connection().close()
-signals.request_finished.connect(_close_connection)
+def cursor():
+ """Return a cursor on the readonly database connection."""
+ return connection().cursor()
diff --git a/frontend/settings.py b/frontend/settings.py
index 0e8f502..95ca552 100644
--- a/frontend/settings.py
+++ b/frontend/settings.py
@@ -22,10 +22,13 @@
MANAGERS = ADMINS
AUTOTEST_DEFAULT = database_settings_helper.get_default_db_config()
+AUTOTEST_GLOBAL = database_settings_helper.get_global_db_config()
+AUTOTEST_READONLY = database_settings_helper.get_readonly_db_config()
ALLOWED_HOSTS = '*'
-DATABASES = {'default': AUTOTEST_DEFAULT,}
+DATABASES = {'default': AUTOTEST_DEFAULT,
+ 'readonly': AUTOTEST_READONLY,}
# Have to set SECRET_KEY before importing connections because of this bug:
# https://code.djangoproject.com/ticket/20704
diff --git a/frontend/setup_django_lite_environment.py b/frontend/setup_django_lite_environment.py
index a31682a..d9b3906 100644
--- a/frontend/setup_django_lite_environment.py
+++ b/frontend/setup_django_lite_environment.py
@@ -5,7 +5,7 @@
'autotest_lib.frontend.settings_lite')
from autotest_lib.frontend.afe import readonly_connection
-readonly_connection.ReadOnlyConnection.set_globally_disabled(True)
+readonly_connection.set_globally_disabled(True)
import django.core.management
diff --git a/frontend/setup_test_environment.py b/frontend/setup_test_environment.py
index 0ca2788..9fe58f5 100644
--- a/frontend/setup_test_environment.py
+++ b/frontend/setup_test_environment.py
@@ -12,33 +12,44 @@
'autotest_lib.frontend.db.backends.afe_sqlite')
settings.DATABASES['default']['NAME'] = ':memory:'
-from django.db import connection
+settings.DATABASES['readonly'] = {}
+settings.DATABASES['readonly']['ENGINE'] = (
+ 'autotest_lib.frontend.db.backends.afe_sqlite')
+settings.DATABASES['readonly']['NAME'] = ':memory:'
+
+from django.db import connections
from autotest_lib.frontend.afe import readonly_connection
+connection = connections['default']
+connection_readonly = connections['readonly']
+
def run_syncdb(verbosity=0):
management.call_command('syncdb', verbosity=verbosity, interactive=False)
-
+ management.call_command('syncdb', verbosity=verbosity, interactive=False,
+ database='readonly')
def destroy_test_database():
connection.close()
+ connection_readonly.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 move to Django 1.x.
- real_connection = connection.connection
- if real_connection is not None:
- real_connection.close()
- connection.connection = None
+ # when we upgrade Django.
+ for con in [connection, connection_readonly]:
+ real_connection = con.connection
+ if real_connection is not None:
+ real_connection.close()
+ con.connection = None
def set_up():
run_syncdb()
- readonly_connection.ReadOnlyConnection.set_globally_disabled(True)
+ readonly_connection.set_globally_disabled(True)
def tear_down():
- readonly_connection.ReadOnlyConnection.set_globally_disabled(False)
+ readonly_connection.set_globally_disabled(False)
destroy_test_database()
diff --git a/scheduler/scheduler_lib.py b/scheduler/scheduler_lib.py
index 6e72571..cdd6c18 100644
--- a/scheduler/scheduler_lib.py
+++ b/scheduler/scheduler_lib.py
@@ -69,7 +69,7 @@
"""
self.db_connection = None
# bypass the readonly connection
- readonly_connection.ReadOnlyConnection.set_globally_disabled(readonly)
+ readonly_connection.set_globally_disabled(readonly)
if autocommit:
# ensure Django connection is in autocommit
setup_django_environment.enable_autocommit()
diff --git a/scheduler/scheduler_lib_unittest.py b/scheduler/scheduler_lib_unittest.py
index b8e3f84..27d4906 100644
--- a/scheduler/scheduler_lib_unittest.py
+++ b/scheduler/scheduler_lib_unittest.py
@@ -9,7 +9,7 @@
from autotest_lib.database import database_connection
from autotest_lib.frontend import setup_django_environment
-from autotest_lib.frontend.afe.readonly_connection import ReadOnlyConnection
+from autotest_lib.frontend.afe import readonly_connection
from autotest_lib.scheduler import scheduler_lib
from django.db import utils as django_utils
@@ -19,13 +19,13 @@
def setUp(self):
self.connection_manager = None
- ReadOnlyConnection.set_globally_disabled = mock.MagicMock()
+ readonly_connection.set_globally_disabled = mock.MagicMock()
setup_django_environment.enable_autocommit = mock.MagicMock()
scheduler_lib.Singleton._instances = {}
def tearDown(self):
- ReadOnlyConnection.set_globally_disabled.reset_mock()
+ readonly_connection.set_globally_disabled.reset_mock()
setup_django_environment.enable_autocommit.reset_mock()
@@ -68,10 +68,10 @@
"""Test that the singleton works as expected."""
# Confirm that instantiating the class applies global db settings.
connection_manager = scheduler_lib.ConnectionManager()
- ReadOnlyConnection.set_globally_disabled.assert_called_once_with(False)
+ readonly_connection.set_globally_disabled.assert_called_once_with(True)
setup_django_environment.enable_autocommit.assert_called_once_with()
- ReadOnlyConnection.set_globally_disabled.reset_mock()
+ readonly_connection.set_globally_disabled.reset_mock()
setup_django_environment.enable_autocommit.reset_mock()
# Confirm that instantiating another connection manager doesn't change
@@ -79,7 +79,7 @@
connection_manager_2 = scheduler_lib.ConnectionManager()
self.assertTrue(connection_manager == connection_manager_2)
self.assertTrue(
- ReadOnlyConnection.set_globally_disabled.call_count == 0)
+ readonly_connection.set_globally_disabled.call_count == 0)
self.assertTrue(
setup_django_environment.enable_autocommit.call_count == 0)
