1# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""RDB request managers and requests.
6
7RDB request managers: Call an rdb api_method with a list of RDBRequests, and
8match the requests to the responses returned.
9
10RDB Request classes: Used in conjunction with the request managers. Each class
11defines the set of fields the rdb needs to fulfill the request, and a hashable
12request object the request managers use to identify a response with a request.
13"""
14
15import collections
16
17import common
18from autotest_lib.scheduler import rdb_utils
19
20
21class RDBRequestManager(object):
22    """Base request manager for RDB requests.
23
24    Each instance of a request manager is associated with one request, and
25    one api call. All subclasses maintain a queue of unexecuted requests, and
26    and expose an api to add requests/retrieve the response for these requests.
27    """
28
29
30    def __init__(self, request, api_call):
31        """
32        @param request: A subclass of rdb_utls.RDBRequest. The manager can only
33            manage requests of one type.
34        @param api_call: The rdb api call this manager is expected to make.
35            A manager can only send requests of type request, to this api call.
36        """
37        self.request = request
38        self.api_call = api_call
39        self.request_queue = []
40
41
42    def add_request(self, **kwargs):
43        """Add an RDBRequest to the queue."""
44        self.request_queue.append(self.request(**kwargs).get_request())
45
46
47    def response(self):
48        """Execute the api call and return a response for each request.
49
50        The order of responses is the same as the order of requests added
51        to the queue.
52
53        @yield: A response for each request added to the queue after the
54            last invocation of response.
55        """
56        if not self.request_queue:
57            raise rdb_utils.RDBException('No requests. Call add_requests '
58                    'with the appropriate kwargs, before calling response.')
59
60        result = self.api_call(self.request_queue)
61        requests = self.request_queue
62        self.request_queue = []
63        for request in requests:
64            yield result.get(request) if result else None
65
66
67class BaseHostRequestManager(RDBRequestManager):
68    """Manager for batched get requests on hosts."""
69
70
71    def response(self):
72        """Yields a popped host from the returned host list."""
73
74        # As a side-effect of returning a host, this method also removes it
75        # from the list of hosts matched up against a request. Eg:
76        #    hqes: [hqe1, hqe2, hqe3]
77        #    client requests: [c_r1, c_r2, c_r3]
78        #    generate requests in rdb: [r1 (c_r1 and c_r2), r2]
79        #    and response {r1: [h1, h2], r2:[h3]}
80        # c_r1 and c_r2 need to get different hosts though they're the same
81        # request, because they're from different queue_entries.
82        for hosts in super(BaseHostRequestManager, self).response():
83            yield hosts.pop() if hosts else None
84
85
86class RDBRequestMeta(type):
87    """Metaclass for constructing rdb requests.
88
89    This meta class creates a read-only request template by combining the
90    request_arguments of all classes in the inheritence hierarchy into a
91    namedtuple.
92    """
93    def __new__(cls, name, bases, dctn):
94        for base in bases:
95            try:
96                dctn['_request_args'].update(base._request_args)
97            except AttributeError:
98                pass
99        dctn['template'] = collections.namedtuple('template',
100                                                  dctn['_request_args'])
101        return type.__new__(cls, name, bases, dctn)
102
103
104class RDBRequest(object):
105    """Base class for an rdb request.
106
107    All classes inheriting from RDBRequest will need to specify a list of
108    request_args necessary to create the request, and will in turn get a
109    request that the rdb understands.
110    """
111    __metaclass__ = RDBRequestMeta
112    __slots__ = set(['_request_args', '_request'])
113    _request_args = set([])
114
115
116    def __init__(self, **kwargs):
117        for key,value in kwargs.iteritems():
118            try:
119                hash(value)
120            except TypeError as e:
121                raise rdb_utils.RDBException('All fields of a %s must be. '
122                        'hashable %s: %s, %s failed this test.' %
123                        (self.__class__, key, type(value), value))
124        try:
125            self._request = self.template(**kwargs)
126        except TypeError:
127            raise rdb_utils.RDBException('Creating %s requires args %s got %s' %
128                    (self.__class__, self.template._fields, kwargs.keys()))
129
130
131    def get_request(self):
132        """Returns a request that the rdb understands.
133
134        @return: A named tuple with all the fields necessary to make a request.
135        """
136        return self._request
137
138
139class HashableDict(dict):
140    """A hashable dictionary.
141
142    This class assumes all values of the input dict are hashable.
143    """
144
145    def __hash__(self):
146        return hash(tuple(sorted(self.items())))
147
148
149class HostRequest(RDBRequest):
150    """Basic request for information about a single host.
151
152    Eg: HostRequest(host_id=x): Will return all information about host x.
153    """
154    _request_args =  set(['host_id'])
155
156
157class UpdateHostRequest(HostRequest):
158    """Defines requests to update hosts.
159
160    Eg:
161        UpdateHostRequest(host_id=x, payload={'afe_hosts_col_name': value}):
162            Will update column afe_hosts_col_name with the given value, for
163            the given host_id.
164
165    @raises RDBException: If the input arguments don't contain the expected
166        fields to make the request, or are of the wrong type.
167    """
168    _request_args = set(['payload'])
169
170
171    def __init__(self, **kwargs):
172        try:
173            kwargs['payload'] = HashableDict(kwargs['payload'])
174        except (KeyError, TypeError) as e:
175            raise rdb_utils.RDBException('Creating %s requires args %s got %s' %
176                    (self.__class__, self.template._fields, kwargs.keys()))
177        super(UpdateHostRequest, self).__init__(**kwargs)
178
179
180class AcquireHostRequest(HostRequest):
181    """Defines requests to acquire hosts.
182
183    Eg:
184        AcquireHostRequest(host_id=None, deps=[d1, d2], acls=[a1, a2],
185                priority=None, parent_job_id=None): Will acquire and return a
186                host that matches the specified deps/acls.
187        AcquireHostRequest(host_id=x, deps=[d1, d2], acls=[a1, a2]) : Will
188            acquire and return host x, after checking deps/acls match.
189
190    @raises RDBException: If the the input arguments don't contain the expected
191        fields to make a request, or are of the wrong type.
192    """
193    # TODO(beeps): Priority and parent_job_id shouldn't be a part of the
194    # core request.
195    _request_args = set(['priority', 'deps', 'preferred_deps', 'acls',
196                         'parent_job_id', 'suite_min_duts'])
197
198
199    def __init__(self, **kwargs):
200        try:
201            kwargs['deps'] = frozenset(kwargs['deps'])
202            kwargs['preferred_deps'] = frozenset(kwargs['preferred_deps'])
203            kwargs['acls'] = frozenset(kwargs['acls'])
204
205            # parent_job_id defaults to NULL but always serializing it as an int
206            # fits the rdb's type assumptions. Note that job ids are 1 based.
207            if kwargs['parent_job_id'] is None:
208                kwargs['parent_job_id'] = 0
209        except (KeyError, TypeError) as e:
210            raise rdb_utils.RDBException('Creating %s requires args %s got %s' %
211                    (self.__class__, self.template._fields, kwargs.keys()))
212        super(AcquireHostRequest, self).__init__(**kwargs)
213
214
215