scheduler/rdb_utils.py

"""RDB utilities.

Do not import rdb or autotest modules here to avoid cyclic dependencies.
"""
import itertools
import collections


class RDBException(Exception):
    """Generic RDB exception."""

    def wire_format(self):
        """Convert the exception to a format better suited to an rpc response.
        """
        return str(self)


class RDBRequestMeta(type):
    """Metaclass for constructing rdb requests.

    This meta class creates a read-only request template by combining the
    request_arguments of all classes in the inheritence hierarchy into a namedtuple.
    """
    def __new__(cls, name, bases, dctn):
        for base in bases:
            try:
                dctn['_request_args'].update(base._request_args)
            except AttributeError:
                pass
        dctn['template'] = collections.namedtuple('template',
                                                  dctn['_request_args'])
        return type.__new__(cls, name, bases, dctn)


# RDB Request classes: Used in conjunction with the request managers defined in
# rdb_lib. Each class defines the set of fields the rdb needs to fulfill the
# request, and a hashable request object the request managers use to identify
# a response with a request.
class RDBRequest(object):
    """Base class for an rdb request.

    All classes inheriting from RDBRequest will need to specify a list of
    request_args necessary to create the request, and will in turn get a
    request that the rdb understands.
    """
    __metaclass__ = RDBRequestMeta
    __slots__ = set(['_request_args', '_request'])
    _request_args = set([])


    def __init__(self, **kwargs):
        for key,value in kwargs.iteritems():
            try:
                hash(value)
            except TypeError as e:
                raise RDBException('All fields of a %s must be hashable. '
                                   '%s: %s, %s failed this test.' %
                                   (self.__class__, key, type(value), value))
        try:
            self._request = self.template(**kwargs)
        except TypeError:
            raise RDBException('Creating %s requires args %s, got %s ' %
                    (self.__class__, self.template._fields, kwargs.keys()))


    def get_request(self):
        """Returns a request that the rdb understands.

        @return: A named tuple with all the fields necessary to make a request.
        """
        return self._request


class HashableDict(dict):
    """A hashable dictionary.

    This class assumes all values of the input dict are hashable.
    """

    def __hash__(self):
        return hash(tuple(sorted(self.items())))


class HostRequest(RDBRequest):
    """Basic request for information about a single host.

    Eg: HostRequest(host_id=x): Will return all information about host x.
    """
    _request_args =  set(['host_id'])


class UpdateHostRequest(HostRequest):
    """Defines requests to update hosts.

    Eg:
        UpdateHostRequest(host_id=x, payload={'afe_hosts_col_name': value}):
            Will update column afe_hosts_col_name with the given value, for
            the given host_id.

    @raises RDBException: If the input arguments don't contain the expected
        fields to make the request, or are of the wrong type.
    """
    _request_args = set(['payload'])


    def __init__(self, **kwargs):
        try:
            kwargs['payload'] = HashableDict(kwargs['payload'])
        except (KeyError, TypeError) as e:
            raise RDBException('Creating %s requires args %s, got %s ' %
                    (self.__class__, self.template._fields, kwargs.keys()))
        super(UpdateHostRequest, self).__init__(**kwargs)


class AcquireHostRequest(HostRequest):
    """Defines requests to acquire hosts.

    Eg:
        AcquireHostRequest(host_id=None, deps=[d1, d2], acls=[a1], priority=40)
            Will acquire and return a host that matches the specified deps/acls
            after doing so for all requests of higher priority first.
        AcquireHostRequest(host_id=x, deps=[d1, d2], acls=[a1, a2], priority=40)
            Will acquire and return host x, if the there are any left after
            satisfying requests of a higher priority.

    @raises RDBException: If the the input arguments don't contain the expected
        fields to make a request, or are of the wrong type.
    """
    _request_args = set(['deps', 'acls', 'priority'])


    def __init__(self, **kwargs):
        try:
            kwargs['deps'] = frozenset(kwargs['deps'])
            kwargs['acls'] = frozenset(kwargs['acls'])
        except (KeyError, TypeError) as e:
            raise RDBException('Creating %s requires args %s, got %s ' %
                    (self.__class__, self.template._fields, kwargs.keys()))
        super(AcquireHostRequest, self).__init__(**kwargs)


# Custom iterators: Used by the rdb to lazily convert the iteration of a
# queryset to a database query and return an appropriately formatted response.
class RememberingIterator(object):
    """An iterator capable of reproducing all values in the input generator.
    """

    #pylint: disable-msg=C0111
    def __init__(self, gen):
        self.current, self.history = itertools.tee(gen)
        self.items = []


    def __iter__(self):
        return self


    def next(self):
        return self.current.next()


    def get_all_items(self):
        """Get all the items in the generator this object was created with.

        @return: A list of items.
        """
        if not self.items:
            self.items = list(self.history)
        return self.items


class LabelIterator(RememberingIterator):
    """A RememberingIterator for labels.

    Within the rdb any label/dependency comparisons are performed based on label
    ids. However, the host object returned needs to contain label names instead.
    This class returns the label id when iterated over, but a list of all label
    names when accessed through get_all_items.
    """


    def next(self):
        return super(LabelIterator, self).next().id


    def get_all_items(self):
        """Get all label names of the labels in the input generator.

        @return: A list of label names.
        """
        return [label.name
                for label in super(LabelIterator, self).get_all_items()]


# Rdb host adapters: Help in making a raw database host object more ameanable
# to the classes and functions in the rdb and/or rdb clients.
class RDBServerHostWrapper(object):
    """A host wrapper for the raw database object.
    """


    def __init__(self, host):
        self.id = host.id
        self.hostname = host.hostname
        self.status = host.status
        self.protection = host.protection
        self.dirty = host.dirty
        self.invalid = host.invalid
        self.labels = LabelIterator(
                (label for label in host.labels.all()))
        self.acls = RememberingIterator(
                (acl.id for acl in host.aclgroup_set.all()))
        platform = host.platform()
        self.platform = platform.name if platform else None


    def wire_format(self):
        """Returns all information needed to scheduler jobs on the host.

        @return: A dictionary of host information.
        """
        host_info = {}
        for key, value in self.__dict__.iteritems():
            if isinstance(value, RememberingIterator):
                host_info[key] = value.get_all_items()
            else:
                host_info[key] = value
        return host_info


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
        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 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