1489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B# Copyright (c) 2014 The Chromium OS Authors. All rights reserved. 2489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B# Use of this source code is governed by a BSD-style license that can be 3489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B# found in the LICENSE file. 4489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 5489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B"""RDB request managers and requests. 6489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 7489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth BRDB request managers: Call an rdb api_method with a list of RDBRequests, and 8489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bmatch the requests to the responses returned. 9489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 10489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth BRDB Request classes: Used in conjunction with the request managers. Each class 11489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bdefines the set of fields the rdb needs to fulfill the request, and a hashable 12489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Brequest object the request managers use to identify a response with a request. 13489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B""" 14489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 15489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bimport collections 16489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 17489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bimport common 18489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bfrom autotest_lib.scheduler import rdb_utils 19489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 20489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 21489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bclass RDBRequestManager(object): 22489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Base request manager for RDB requests. 23489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 24489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B Each instance of a request manager is associated with one request, and 25489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B one api call. All subclasses maintain a queue of unexecuted requests, and 26489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B and expose an api to add requests/retrieve the response for these requests. 27489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 28489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 29489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 30489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def __init__(self, request, api_call): 31489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 32489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B @param request: A subclass of rdb_utls.RDBRequest. The manager can only 33489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B manage requests of one type. 34489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B @param api_call: The rdb api call this manager is expected to make. 35489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B A manager can only send requests of type request, to this api call. 36489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 37489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B self.request = request 38489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B self.api_call = api_call 39489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B self.request_queue = [] 40489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 41489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 42489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def add_request(self, **kwargs): 43489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Add an RDBRequest to the queue.""" 44489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B self.request_queue.append(self.request(**kwargs).get_request()) 45489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 46489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 47489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def response(self): 48489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Execute the api call and return a response for each request. 49489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 50489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B The order of responses is the same as the order of requests added 51489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B to the queue. 52489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 53489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B @yield: A response for each request added to the queue after the 54489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B last invocation of response. 55489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 56489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B if not self.request_queue: 57489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B raise rdb_utils.RDBException('No requests. Call add_requests ' 58489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 'with the appropriate kwargs, before calling response.') 59489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 60489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B result = self.api_call(self.request_queue) 61489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B requests = self.request_queue 62489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B self.request_queue = [] 63489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B for request in requests: 64489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B yield result.get(request) if result else None 65489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 66489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 67489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bclass BaseHostRequestManager(RDBRequestManager): 68489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Manager for batched get requests on hosts.""" 69489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 70489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 71489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def response(self): 72489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Yields a popped host from the returned host list.""" 73489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 74489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B # As a side-effect of returning a host, this method also removes it 75489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B # from the list of hosts matched up against a request. Eg: 76489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B # hqes: [hqe1, hqe2, hqe3] 77489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B # client requests: [c_r1, c_r2, c_r3] 78489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B # generate requests in rdb: [r1 (c_r1 and c_r2), r2] 79489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B # and response {r1: [h1, h2], r2:[h3]} 80489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B # c_r1 and c_r2 need to get different hosts though they're the same 81489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B # request, because they're from different queue_entries. 82489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B for hosts in super(BaseHostRequestManager, self).response(): 83489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B yield hosts.pop() if hosts else None 84489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 85489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 86489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bclass RDBRequestMeta(type): 87489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Metaclass for constructing rdb requests. 88489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 89489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B This meta class creates a read-only request template by combining the 90489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B request_arguments of all classes in the inheritence hierarchy into a 91489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B namedtuple. 92489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 93489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def __new__(cls, name, bases, dctn): 94489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B for base in bases: 95489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B try: 96489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B dctn['_request_args'].update(base._request_args) 97489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B except AttributeError: 98489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B pass 99489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B dctn['template'] = collections.namedtuple('template', 100489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B dctn['_request_args']) 101489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B return type.__new__(cls, name, bases, dctn) 102489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 103489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 104489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bclass RDBRequest(object): 105489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Base class for an rdb request. 106489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 107489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B All classes inheriting from RDBRequest will need to specify a list of 108489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B request_args necessary to create the request, and will in turn get a 109489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B request that the rdb understands. 110489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 111489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B __metaclass__ = RDBRequestMeta 112489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B __slots__ = set(['_request_args', '_request']) 113489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B _request_args = set([]) 114489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 115489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 116489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def __init__(self, **kwargs): 117489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B for key,value in kwargs.iteritems(): 118489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B try: 119489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B hash(value) 120489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B except TypeError as e: 121489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B raise rdb_utils.RDBException('All fields of a %s must be. ' 122489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 'hashable %s: %s, %s failed this test.' % 123489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B (self.__class__, key, type(value), value)) 124489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B try: 125489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B self._request = self.template(**kwargs) 126489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B except TypeError: 127489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B raise rdb_utils.RDBException('Creating %s requires args %s got %s' % 128489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B (self.__class__, self.template._fields, kwargs.keys())) 129489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 130489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 131489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def get_request(self): 132489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Returns a request that the rdb understands. 133489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 134489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B @return: A named tuple with all the fields necessary to make a request. 135489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 136489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B return self._request 137489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 138489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 139489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bclass HashableDict(dict): 140489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """A hashable dictionary. 141489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 142489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B This class assumes all values of the input dict are hashable. 143489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 144489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 145489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def __hash__(self): 146489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B return hash(tuple(sorted(self.items()))) 147489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 148489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 149489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bclass HostRequest(RDBRequest): 150489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Basic request for information about a single host. 151489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 152489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B Eg: HostRequest(host_id=x): Will return all information about host x. 153489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 154489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B _request_args = set(['host_id']) 155489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 156489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 157489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bclass UpdateHostRequest(HostRequest): 158489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Defines requests to update hosts. 159489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 160489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B Eg: 161489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B UpdateHostRequest(host_id=x, payload={'afe_hosts_col_name': value}): 162489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B Will update column afe_hosts_col_name with the given value, for 163489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B the given host_id. 164489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 165489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B @raises RDBException: If the input arguments don't contain the expected 166489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B fields to make the request, or are of the wrong type. 167489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 168489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B _request_args = set(['payload']) 169489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 170489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 171489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def __init__(self, **kwargs): 172489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B try: 173489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B kwargs['payload'] = HashableDict(kwargs['payload']) 174489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B except (KeyError, TypeError) as e: 175489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B raise rdb_utils.RDBException('Creating %s requires args %s got %s' % 176489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B (self.__class__, self.template._fields, kwargs.keys())) 177489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B super(UpdateHostRequest, self).__init__(**kwargs) 178489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 179489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 180489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth Bclass AcquireHostRequest(HostRequest): 181489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """Defines requests to acquire hosts. 182489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 183489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B Eg: 1842c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B AcquireHostRequest(host_id=None, deps=[d1, d2], acls=[a1, a2], 1852c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B priority=None, parent_job_id=None): Will acquire and return a 1862c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B host that matches the specified deps/acls. 187489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B AcquireHostRequest(host_id=x, deps=[d1, d2], acls=[a1, a2]) : Will 188489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B acquire and return host x, after checking deps/acls match. 189489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 190489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B @raises RDBException: If the the input arguments don't contain the expected 191489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B fields to make a request, or are of the wrong type. 192489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B """ 1932c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B # TODO(beeps): Priority and parent_job_id shouldn't be a part of the 1942c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B # core request. 195a9bc9592e64c5fc6088e07ab07826a0d5762ba54Fang Deng _request_args = set(['priority', 'deps', 'preferred_deps', 'acls', 196a9bc9592e64c5fc6088e07ab07826a0d5762ba54Fang Deng 'parent_job_id', 'suite_min_duts']) 197489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 198489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 199489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B def __init__(self, **kwargs): 200489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B try: 201489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B kwargs['deps'] = frozenset(kwargs['deps']) 202a9bc9592e64c5fc6088e07ab07826a0d5762ba54Fang Deng kwargs['preferred_deps'] = frozenset(kwargs['preferred_deps']) 203489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B kwargs['acls'] = frozenset(kwargs['acls']) 2042c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B 2052c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B # parent_job_id defaults to NULL but always serializing it as an int 2062c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B # fits the rdb's type assumptions. Note that job ids are 1 based. 2072c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B if kwargs['parent_job_id'] is None: 2082c1a22a9f93bf50147cd4e6b10487d02768f8919Prashanth B kwargs['parent_job_id'] = 0 209489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B except (KeyError, TypeError) as e: 210489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B raise rdb_utils.RDBException('Creating %s requires args %s got %s' % 211489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B (self.__class__, self.template._fields, kwargs.keys())) 212489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B super(AcquireHostRequest, self).__init__(**kwargs) 213489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 214489b91d72cd225e902081dbd3f9e47448fe867f6Prashanth B 215