def _ensure_label_exists(name):
"""
Ensure that a label called |name| exists in the Django models.
This function is to be called from within afe rpcs only, as an
alternative to server.cros.provision.ensure_label_exists(...). It works
by Django model manipulation, rather than by making another create_label
rpc call.
@param name: the label to check for/create.
@raises ValidationError: There was an error in the response that was
not because the label already existed.
@returns True is a label was created, False otherwise.
"""
# Make sure this function is not called on shards but only on master.
assert not server_utils.is_shard()
try:
models.Label.objects.get(name=name)
except models.Label.DoesNotExist:
try:
new_label = models.Label.objects.create(name=name)
new_label.save()
return True
except django.db.utils.IntegrityError as e:
# It is possible that another suite/test already
# created the label between the check and save.
if DUPLICATE_KEY_MSG in str(e):
return False
else:
raise
return False
def should_cleanup():
"""Check if the server should run db_cleanup.
Only shard should clean up db.
@returns: True if it should run db cleanup otherwise False.
"""
return utils.is_shard()
def _validate_host_job_sharding(host_objects):
"""Check that the hosts obey job sharding rules."""
if not (server_utils.is_shard()
or _allowed_hosts_for_master_job(host_objects)):
shard_host_map = bucket_hosts_by_shard(host_objects)
raise ValueError('The following hosts are on shard(s), please create '
'seperate jobs for hosts on each shard: %s ' %
shard_host_map)
def should_optimize():
"""Check if the server should run db_optimize.
Only shard should optimize db.
@returns: True if it should optimize db otherwise False.
"""
return utils.is_shard()
def run_rpc_on_multiple_hostnames(rpc_call, shard_hostnames, **kwargs):
"""Runs an rpc to multiple AFEs
This is i.e. used to propagate changes made to hosts after they are assigned
to a shard.
@param rpc_call: Name of the rpc endpoint to call.
@param shard_hostnames: List of hostnames to run the rpcs on.
@param **kwargs: Keyword arguments to pass in the rpcs.
"""
# Make sure this function is not called on shards but only on master.
assert not server_utils.is_shard()
for shard_hostname in shard_hostnames:
afe = frontend_wrappers.RetryingAFE(server=shard_hostname,
user=thread_local.get_user())
afe.run(rpc_call, **kwargs)
def replacement(*args, **kwargs):
"""We need special handling when decorating an RPC that can be called
directly using positional arguments.
One example is rpc_interface.create_job().
rpc_interface.create_job_page_handler() calls the function using both
positional and keyword arguments. Since frontend.RpcClient.run()
takes only keyword arguments for an RPC, positional arguments of the
RPC function need to be transformed into keyword arguments.
"""
kwargs = _convert_to_kwargs_only(func, args, kwargs)
if server_utils.is_shard():
afe = frontend_wrappers.RetryingAFE(
server=server_utils.get_global_afe_hostname(),
user=thread_local.get_user())
return afe.run(func.func_name, **kwargs)
return func(**kwargs)
def replacement(**kwargs):
# Only keyword arguments can be accepted here, as we need the argument
# names to send the rpc. serviceHandler always provides arguments with
# their keywords, so this is not a problem.
# A host record (identified by kwargs['id']) can be deleted in
# func(). Therefore, we should save the data that can be needed later
# before func() is called.
shard_hostname = None
host = models.Host.smart_get(kwargs['id'])
if host and host.shard:
shard_hostname = host.shard.hostname
ret = func(**kwargs)
if shard_hostname and not server_utils.is_shard():
run_rpc_on_multiple_hostnames(func.func_name, [shard_hostname],
**kwargs)
return ret
from django.db.models import fields
from django.db.models import Q
from autotest_lib.frontend.afe import models
from autotest_lib.scheduler import rdb_cache_manager
from autotest_lib.scheduler import rdb_hosts
from autotest_lib.scheduler import rdb_requests
from autotest_lib.scheduler import rdb_utils
from autotest_lib.server import utils
try:
from chromite.lib import metrics
except ImportError:
metrics = utils.metrics_mock
_rdb_timer_name = 'chromeos/autotest/scheduler/rdb/durations/%s'
_is_master = not utils.is_shard()
# Qeury managers: Provide a layer of abstraction over the database by
# encapsulating common query patterns used by the rdb.
class BaseHostQueryManager(object):
"""Base manager for host queries on all hosts.
"""
host_objects = models.Host.objects
def update_hosts(self, host_ids, **kwargs):
"""Update fields on a hosts.
@param host_ids: A list of ids of hosts to update.
@param kwargs: A key value dictionary corresponding to column, value