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 Host objects.
6
7RDBHost: Basic host object, capable of retrieving fields of a host that
8correspond to columns of the host table.
9
10RDBServerHostWrapper: Server side host adapters that help in making a raw
11database host object more ameanable to the classes and functions in the rdb
12and/or rdb clients.
13
14RDBClientHostWrapper: Scheduler host proxy that converts host information
15returned by the rdb into a client host object capable of proxying updates
16back to the rdb.
17"""
18
19import logging
20import time
21from django.core import exceptions as django_exceptions
22
23import common
24from autotest_lib.frontend.afe import rdb_model_extensions as rdb_models
25from autotest_lib.frontend.afe import models as afe_models
26from autotest_lib.scheduler import rdb_requests
27from autotest_lib.scheduler import rdb_utils
28from autotest_lib.site_utils import metadata_reporter
29from autotest_lib.site_utils.suite_scheduler import constants
30
31
32class RDBHost(object):
33    """A python host object representing a django model for the host."""
34
35    required_fields = set(
36            rdb_models.AbstractHostModel.get_basic_field_names() + ['id'])
37
38
39    def _update_attributes(self, new_attributes):
40        """Updates attributes based on an input dictionary.
41
42        Since reads are not proxied to the rdb this method caches updates to
43        the host tables as class attributes.
44
45        @param new_attributes: A dictionary of attributes to update.
46        """
47        for name, value in new_attributes.iteritems():
48            setattr(self, name, value)
49
50
51    def __init__(self, **kwargs):
52        if self.required_fields - set(kwargs.keys()):
53            raise rdb_utils.RDBException('Creating %s requires %s, got %s '
54                    % (self.__class__, self.required_fields, kwargs.keys()))
55        self._update_attributes(kwargs)
56
57
58    @classmethod
59    def get_required_fields_from_host(cls, host):
60        """Returns all required attributes of the host parsed into a dict.
61
62        Required attributes are defined as the attributes required to
63        create an RDBHost, and mirror the columns of the host table.
64
65        @param host: A host object containing all required fields as attributes.
66        """
67        required_fields_map = {}
68        try:
69            for field in cls.required_fields:
70                required_fields_map[field] = getattr(host, field)
71        except AttributeError as e:
72            raise rdb_utils.RDBException('Required %s' % e)
73        required_fields_map['id'] = host.id
74        return required_fields_map
75
76
77    def wire_format(self):
78        """Returns information about this host object.
79
80        @return: A dictionary of fields representing the host.
81        """
82        return RDBHost.get_required_fields_from_host(self)
83
84
85class RDBServerHostWrapper(RDBHost):
86    """A host wrapper for the base host object.
87
88    This object contains all the attributes of the raw database columns,
89    and a few more that make the task of host assignment easier. It handles
90    the following duties:
91        1. Serialization of the host object and foreign keys
92        2. Conversion of label ids to label names, and retrieval of platform
93        3. Checking the leased bit/status of a host before leasing it out.
94    """
95
96    def __init__(self, host):
97        """Create an RDBServerHostWrapper.
98
99        @param host: An instance of the Host model class.
100        """
101        host_fields = RDBHost.get_required_fields_from_host(host)
102        super(RDBServerHostWrapper, self).__init__(**host_fields)
103        self.labels = rdb_utils.LabelIterator(host.labels.all())
104        self.acls = [aclgroup.id for aclgroup in host.aclgroup_set.all()]
105        self.protection = host.protection
106        platform = host.platform()
107        # Platform needs to be a method, not an attribute, for
108        # backwards compatibility with the rest of the host model.
109        self.platform_name = platform.name if platform else None
110        self.shard_id = host.shard_id
111
112
113    def refresh(self, fields=None):
114        """Refresh the attributes on this instance.
115
116        @param fields: A list of fieldnames to refresh. If None
117            all the required fields of the host are refreshed.
118
119        @raises RDBException: If refreshing a field fails.
120        """
121        # TODO: This is mainly required for cache correctness. If it turns
122        # into a bottleneck, cache host_ids instead of rdbhosts and rebuild
123        # the hosts once before leasing them out. The important part is to not
124        # trust the leased bit on a cached host.
125        fields = self.required_fields if not fields else fields
126        try:
127            refreshed_fields = afe_models.Host.objects.filter(
128                    id=self.id).values(*fields)[0]
129        except django_exceptions.FieldError as e:
130            raise rdb_utils.RDBException('Couldn\'t refresh fields %s: %s' %
131                    fields, e)
132        self._update_attributes(refreshed_fields)
133
134
135    def lease(self):
136        """Set the leased bit on the host object, and in the database.
137
138        @raises RDBException: If the host is already leased.
139        """
140        self.refresh(fields=['leased'])
141        if self.leased:
142            raise rdb_utils.RDBException('Host %s is already leased' %
143                                         self.hostname)
144        self.leased = True
145        # TODO: Avoid leaking django out of rdb.QueryManagers. This is still
146        # preferable to calling save() on the host object because we're only
147        # updating/refreshing a single indexed attribute, the leased bit.
148        afe_models.Host.objects.filter(id=self.id).update(leased=self.leased)
149
150
151    def wire_format(self, unwrap_foreign_keys=True):
152        """Returns all information needed to scheduler jobs on the host.
153
154        @param unwrap_foreign_keys: If true this method will retrieve and
155            serialize foreign keys of the original host, which are stored
156            in the RDBServerHostWrapper as iterators.
157
158        @return: A dictionary of host information.
159        """
160        host_info = super(RDBServerHostWrapper, self).wire_format()
161
162        if unwrap_foreign_keys:
163            host_info['labels'] = self.labels.get_label_names()
164            host_info['acls'] = self.acls
165            host_info['platform_name'] = self.platform_name
166            host_info['protection'] = self.protection
167        return host_info
168
169
170class RDBClientHostWrapper(RDBHost):
171    """A client host wrapper for the base host object.
172
173    This wrapper is used whenever the queue entry needs direct access
174    to the host.
175    """
176
177    def __init__(self, **kwargs):
178
179        # This class is designed to only check for the bare minimum
180        # attributes on a host, so if a client tries accessing an
181        # unpopulated foreign key it will result in an exception. Doing
182        # so makes it easier to add fields to the rdb host without
183        # updating all the clients.
184        super(RDBClientHostWrapper, self).__init__(**kwargs)
185
186        # TODO(beeps): Remove this once we transition to urls
187        from autotest_lib.scheduler import rdb
188        self.update_request_manager = rdb_requests.RDBRequestManager(
189                rdb_requests.UpdateHostRequest, rdb.update_hosts)
190        self.dbg_str = ''
191        self.metadata = {}
192
193
194    def _update(self, payload):
195        """Send an update to rdb, save the attributes of the payload locally.
196
197        @param: A dictionary representing 'key':value of the update required.
198
199        @raises RDBException: If the update fails.
200        """
201        logging.info('Host %s in %s updating %s through rdb on behalf of: %s ',
202                     self.hostname, self.status, payload, self.dbg_str)
203        self.update_request_manager.add_request(host_id=self.id,
204                payload=payload)
205        for response in self.update_request_manager.response():
206            if response:
207                raise rdb_utils.RDBException('Host %s unable to perform update '
208                        '%s through rdb on behalf of %s: %s',  self.hostname,
209                        payload, self.dbg_str, response)
210        super(RDBClientHostWrapper, self)._update_attributes(payload)
211
212
213    def record_state(self, type_str, state, value):
214        """Record metadata in elasticsearch.
215
216        @param type_str: sets the _type field in elasticsearch db.
217        @param state: string representing what state we are recording,
218                      e.g. 'status'
219        @param value: value of the state, e.g. 'running'
220        """
221        metadata = {
222            state: value,
223            'hostname': self.hostname,
224            'board': self.board,
225            'pools': self.pools,
226            'dbg_str': self.dbg_str,
227            '_type': type_str,
228            'time_recorded': time.time(),
229        }
230        metadata.update(self.metadata)
231        metadata_reporter.queue(metadata)
232
233
234    def set_status(self, status):
235        """Proxy for setting the status of a host via the rdb.
236
237        @param status: The new status.
238        """
239        self._update({'status': status})
240        self.record_state('host_history', 'status', status)
241
242
243    def update_field(self, fieldname, value):
244        """Proxy for updating a field on the host.
245
246        @param fieldname: The fieldname as a string.
247        @param value: The value to assign to the field.
248        """
249        self._update({fieldname: value})
250
251
252    def platform_and_labels(self):
253        """Get the platform and labels on this host.
254
255        @return: A tuple containing a list of label names and the platform name.
256        """
257        platform = self.platform_name
258        labels = [label for label in self.labels if label != platform]
259        return platform, labels
260
261
262    def platform(self):
263        """Get the name of the platform of this host.
264
265        @return: A string representing the name of the platform.
266        """
267        return self.platform_name
268
269
270    def find_labels_start_with(self, search_string):
271        """Find all labels started with given string.
272
273        @param search_string: A string to match the beginning of the label.
274        @return: A list of all matched labels.
275        """
276        try:
277            return [l for l in self.labels if l.startswith(search_string)]
278        except AttributeError:
279            return []
280
281
282    @property
283    def board(self):
284        """Get the names of the board of this host.
285
286        @return: A string of the name of the board, e.g., lumpy.
287        """
288        boards = self.find_labels_start_with(constants.Labels.BOARD_PREFIX)
289        return (boards[0][len(constants.Labels.BOARD_PREFIX):] if boards
290                else None)
291
292
293    @property
294    def pools(self):
295        """Get the names of the pools of this host.
296
297        @return: A list of pool names that the host is assigned to.
298        """
299        return [label[len(constants.Labels.POOL_PREFIX):] for label in
300                self.find_labels_start_with(constants.Labels.POOL_PREFIX)]
301
302
303    def get_object_dict(self, **kwargs):
304        """Serialize the attributes of this object into a dict.
305
306        This method is called through frontend code to get a serialized
307        version of this object.
308
309        @param kwargs:
310            extra_fields: Extra fields, outside the columns of a host table.
311
312        @return: A dictionary representing the fields of this host object.
313        """
314        # TODO(beeps): Implement support for extra fields. Currently nothing
315        # requires them.
316        return self.wire_format()
317
318
319    def save(self):
320        """Save any local data a client of this host object might have saved.
321
322        Setting attributes on a model before calling its save() method is a
323        common django pattern. Most, if not all updates to the host happen
324        either through set status or update_field. Though we keep the internal
325        state of the RDBClientHostWrapper consistent through these updates
326        we need a bulk save method such as this one to save any attributes of
327        this host another model might have set on it before calling its own
328        save method. Eg:
329            task = ST.objects.get(id=12)
330            task.host.status = 'Running'
331            task.save() -> this should result in the hosts status changing to
332            Running.
333
334        Functions like add_host_to_labels will have to update this host object
335        differently, as that is another level of foreign key indirection.
336        """
337        self._update(self.get_required_fields_from_host(self))
338
339
340def return_rdb_host(func):
341    """Decorator for functions that return a list of Host objects.
342
343    @param func: The decorated function.
344    @return: A functions capable of converting each host_object to a
345        rdb_hosts.RDBServerHostWrapper.
346    """
347    def get_rdb_host(*args, **kwargs):
348        """Takes a list of hosts and returns a list of host_infos.
349
350        @param hosts: A list of hosts. Each host is assumed to contain
351            all the fields in a host_info defined above.
352        @return: A list of rdb_hosts.RDBServerHostWrappers, one per host, or an
353            empty list is no hosts were found..
354        """
355        hosts = func(*args, **kwargs)
356        return [RDBServerHostWrapper(host) for host in hosts]
357    return get_rdb_host
358
359
360