| # 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. |
| |
| """RDB Host objects. |
| |
| RDBHost: Basic host object, capable of retrieving fields of a host that |
| correspond to columns of the host table. |
| |
| RDBServerHostWrapper: Server side host adapters that help in making a raw |
| database host object more ameanable to the classes and functions in the rdb |
| and/or rdb clients. |
| |
| RDBClientHostWrapper: Scheduler host proxy that converts host information |
| returned by the rdb into a client host object capable of proxying updates |
| back to the rdb. |
| """ |
| |
| import logging |
| from django.core import exceptions as django_exceptions |
| |
| import common |
| from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models |
| from autotest_lib.frontend.afe import models as afe_models |
| from autotest_lib.scheduler import rdb_requests |
| from autotest_lib.scheduler import rdb_utils |
| |
| |
| class RDBHost(object): |
| """A python host object representing a django model for the host.""" |
| |
| required_fields = set( |
| rdb_models.AbstractHostModel.get_basic_field_names() + ['id']) |
| |
| |
| def _update_attributes(self, new_attributes): |
| """Updates attributes based on an input dictionary. |
| |
| Since reads are not proxied to the rdb this method caches updates to |
| the host tables as class attributes. |
| |
| @param new_attributes: A dictionary of attributes to update. |
| """ |
| for name, value in new_attributes.iteritems(): |
| setattr(self, name, value) |
| |
| |
| def __init__(self, **kwargs): |
| if self.required_fields - set(kwargs.keys()): |
| raise rdb_utils.RDBException('Creating %s requires %s, got %s ' |
| % (self.__class__, self.required_fields, kwargs.keys())) |
| self._update_attributes(kwargs) |
| |
| |
| @classmethod |
| def get_required_fields_from_host(cls, host): |
| """Returns all required attributes of the host parsed into a dict. |
| |
| Required attributes are defined as the attributes required to |
| create an RDBHost, and mirror the columns of the host table. |
| |
| @param host: A host object containing all required fields as attributes. |
| """ |
| required_fields_map = {} |
| try: |
| for field in cls.required_fields: |
| required_fields_map[field] = getattr(host, field) |
| except AttributeError as e: |
| raise rdb_utils.RDBException('Required %s' % e) |
| required_fields_map['id'] = host.id |
| return required_fields_map |
| |
| |
| def wire_format(self): |
| """Returns information about this host object. |
| |
| @return: A dictionary of fields representing the host. |
| """ |
| return RDBHost.get_required_fields_from_host(self) |
| |
| |
| class RDBServerHostWrapper(RDBHost): |
| """A host wrapper for the base host object. |
| |
| This object contains all the attributes of the raw database columns, |
| and a few more that make the task of host assignment easier. It handles |
| the following duties: |
| 1. Serialization of the host object and foreign keys |
| 2. Conversion of label ids to label names, and retrieval of platform |
| 3. Checking the leased bit/status of a host before leasing it out. |
| """ |
| |
| def __init__(self, host): |
| """Create an RDBServerHostWrapper. |
| |
| @param host: An instance of the Host model class. |
| """ |
| host_fields = RDBHost.get_required_fields_from_host(host) |
| super(RDBServerHostWrapper, self).__init__(**host_fields) |
| self.labels = rdb_utils.LabelIterator(host.labels.all()) |
| self.acls = [aclgroup.id for aclgroup in host.aclgroup_set.all()] |
| self.protection = host.protection |
| platform = host.platform() |
| # Platform needs to be a method, not an attribute, for |
| # backwards compatibility with the rest of the host model. |
| self.platform_name = platform.name if platform else None |
| |
| |
| def refresh(self, fields=None): |
| """Refresh the attributes on this instance. |
| |
| @param fields: A list of fieldnames to refresh. If None |
| all the required fields of the host are refreshed. |
| |
| @raises RDBException: If refreshing a field fails. |
| """ |
| # TODO: This is mainly required for cache correctness. If it turns |
| # into a bottleneck, cache host_ids instead of rdbhosts and rebuild |
| # the hosts once before leasing them out. The important part is to not |
| # trust the leased bit on a cached host. |
| fields = self.required_fields if not fields else fields |
| try: |
| refreshed_fields = afe_models.Host.objects.filter( |
| id=self.id).values(*fields)[0] |
| except django_exceptions.FieldError as e: |
| raise rdb_utils.RDBException('Couldn\'t refresh fields %s: %s' % |
| fields, e) |
| self._update_attributes(refreshed_fields) |
| |
| |
| def lease(self): |
| """Set the leased bit on the host object, and in the database. |
| |
| @raises RDBException: If the host is already leased. |
| """ |
| self.refresh(fields=['leased']) |
| if self.leased: |
| raise rdb_utils.RDBException('Host %s is already leased' % |
| self.hostname) |
| self.leased = True |
| # TODO: Avoid leaking django out of rdb.QueryManagers. This is still |
| # preferable to calling save() on the host object because we're only |
| # updating/refreshing a single indexed attribute, the leased bit. |
| afe_models.Host.objects.filter(id=self.id).update(leased=self.leased) |
| |
| |
| def wire_format(self, unwrap_foreign_keys=True): |
| """Returns all information needed to scheduler jobs on the host. |
| |
| @param unwrap_foreign_keys: If true this method will retrieve and |
| serialize foreign keys of the original host, which are stored |
| in the RDBServerHostWrapper as iterators. |
| |
| @return: A dictionary of host information. |
| """ |
| host_info = super(RDBServerHostWrapper, self).wire_format() |
| |
| if unwrap_foreign_keys: |
| host_info['labels'] = self.labels.get_label_names() |
| host_info['acls'] = self.acls |
| host_info['platform_name'] = self.platform_name |
| host_info['protection'] = self.protection |
| return host_info |
| |
| |
| class RDBClientHostWrapper(RDBHost): |
| """A client host wrapper for the base host object. |
| |
| This wrapper is used whenever the queue entry needs direct access |
| to the host. |
| """ |
| |
| def __init__(self, **kwargs): |
| |
| # This class is designed to only check for the bare minimum |
| # attributes on a host, so if a client tries accessing an |
| # unpopulated foreign key it will result in an exception. Doing |
| # so makes it easier to add fields to the rdb host without |
| # updating all the clients. |
| super(RDBClientHostWrapper, self).__init__(**kwargs) |
| |
| # TODO(beeps): Remove this once we transition to urls |
| from autotest_lib.scheduler import rdb |
| self.update_request_manager = rdb_requests.RDBRequestManager( |
| rdb_requests.UpdateHostRequest, rdb.update_hosts) |
| self.dbg_str = '' |
| |
| |
| def _update(self, payload): |
| """Send an update to rdb, save the attributes of the payload locally. |
| |
| @param: A dictionary representing 'key':value of the update required. |
| |
| @raises RDBException: If the update fails. |
| """ |
| logging.info('Host %s in %s updating %s through rdb on behalf of: %s ', |
| self.hostname, self.status, payload, self.dbg_str) |
| self.update_request_manager.add_request(host_id=self.id, |
| payload=payload) |
| for response in self.update_request_manager.response(): |
| if response: |
| raise rdb_utils.RDBException('Host %s unable to perform update ' |
| '%s through rdb on behalf of %s: %s', self.hostname, |
| payload, self.dbg_str, response) |
| super(RDBClientHostWrapper, self)._update_attributes(payload) |
| |
| |
| def set_status(self, status): |
| """Proxy for setting the status of a host via the rdb. |
| |
| @param status: The new status. |
| """ |
| self._update({'status': status}) |
| |
| |
| def update_field(self, fieldname, value): |
| """Proxy for updating a field on the host. |
| |
| @param fieldname: The fieldname as a string. |
| @param value: The value to assign to the field. |
| """ |
| self._update({fieldname: value}) |
| |
| |
| def platform_and_labels(self): |
| """Get the platform and labels on this host. |
| |
| @return: A tuple containing a list of label names and the platform name. |
| """ |
| platform = self.platform_name |
| labels = [label for label in self.labels if label != platform] |
| return platform, labels |
| |
| |
| def platform(self): |
| """Get the name of the platform of this host. |
| |
| @return: A string representing the name of the platform. |
| """ |
| return self.platform_name |
| |
| |
| def get_object_dict(self, **kwargs): |
| """Serialize the attributes of this object into a dict. |
| |
| This method is called through frontend code to get a serialized |
| version of this object. |
| |
| @param kwargs: |
| extra_fields: Extra fields, outside the columns of a host table. |
| |
| @return: A dictionary representing the fields of this host object. |
| """ |
| # TODO(beeps): Implement support for extra fields. Currently nothing |
| # requires them. |
| return self.wire_format() |
| |
| |
| def save(self): |
| """Save any local data a client of this host object might have saved. |
| |
| Setting attributes on a model before calling its save() method is a |
| common django pattern. Most, if not all updates to the host happen |
| either through set status or update_field. Though we keep the internal |
| state of the RDBClientHostWrapper consistent through these updates |
| we need a bulk save method such as this one to save any attributes of |
| this host another model might have set on it before calling its own |
| save method. Eg: |
| task = ST.objects.get(id=12) |
| task.host.status = 'Running' |
| task.save() -> this should result in the hosts status changing to |
| Running. |
| |
| Functions like add_host_to_labels will have to update this host object |
| differently, as that is another level of foreign key indirection. |
| """ |
| self._update(self.get_required_fields_from_host(self)) |
| |
| |
| def return_rdb_host(func): |
| """Decorator for functions that return a list of Host objects. |
| |
| @param func: The decorated function. |
| @return: A functions capable of converting each host_object to a |
| rdb_hosts.RDBServerHostWrapper. |
| """ |
| def get_rdb_host(*args, **kwargs): |
| """Takes a list of hosts and returns a list of host_infos. |
| |
| @param hosts: A list of hosts. Each host is assumed to contain |
| all the fields in a host_info defined above. |
| @return: A list of rdb_hosts.RDBServerHostWrappers, one per host, or an |
| empty list is no hosts were found.. |
| """ |
| hosts = func(*args, **kwargs) |
| return [RDBServerHostWrapper(host) for host in hosts] |
| return get_rdb_host |
| |
| |