1#!/usr/bin/env python
2
3# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7"""
8This script provides functions to:
91. collect: Collect all hosts and their labels to metaDB, can be scheduled
10            run daily, e.g.,
11            ./site_utils/host_label_utils.py collect
122. query: Query for hosts and their labels information at a given day, e.g.,
13          ./site_utils/host_label_utils.py query -n 172.27.213.193 -l peppy
14"""
15
16import argparse
17import itertools
18import logging
19import pprint
20import time
21
22import common
23from autotest_lib.client.common_lib import time_utils
24from autotest_lib.client.common_lib.cros.graphite import autotest_es
25from autotest_lib.frontend import setup_django_environment
26from autotest_lib.frontend.afe import models
27
28
29# _type used for ES
30_HOST_LABEL_TYPE = 'host_labels'
31_HOST_LABEL_TIME_INDEX_TYPE = 'host_labels_time_index'
32
33def get_all_boards(labels=None):
34    """Get a list of boards from host labels.
35
36    Scan through all labels of all duts and get all possible boards based on
37    label of name board:*
38
39    @param labels: A list of labels to filter hosts.
40    @return: A list of board names, e.g., ['peppy', 'daisy']
41    """
42    host_labels = get_host_labels(labels=labels)
43    board_labels = [[label[6:] for label in labels
44                     if label.startswith('board:')]
45                    for labels in host_labels.values()]
46    boards = list(set(itertools.chain.from_iterable(board_labels)))
47    return boards
48
49
50def get_host_labels(days_back=0, hostname=None, labels=None):
51    """Get the labels for a given host or all hosts.
52
53    @param days_back: Get the label info around that number of days back. The
54                      default is 0, i.e., the latest label information.
55    @param hostname: Name of the host, if set to None, return labels for all
56                     hosts. Default is None.
57    @param labels: A list of labels to filter hosts.
58    @return: A dictionary of host labels, key is the hostname, and value is a
59             list of labels, e.g.,
60             {'host1': ['board:daisy', 'pool:bvt']}
61    """
62    # Search for the latest logged labels before the given days_back.
63    # Default is 0, which means the last time host labels were logged.
64    t_end = time.time() - days_back*24*3600
65    results = autotest_es.query(
66            fields_returned=['time_index'],
67            equality_constraints=[('_type', _HOST_LABEL_TIME_INDEX_TYPE),],
68            range_constraints=[('time_index', None, t_end)],
69            size=1,
70            sort_specs=[{'time_index': 'desc'}])
71    t_end_str = time_utils.epoch_time_to_date_string(t_end)
72    if results.total == 0:
73        logging.error('No label information was logged before %s.', t_end_str)
74        return
75    time_index = results.hits[0]['time_index']
76    logging.info('Host labels were recorded at %s',
77                 time_utils.epoch_time_to_date_string(time_index))
78
79    # Search for labels for a given host or all hosts, at time_index.
80    equality_constraints=[('_type', _HOST_LABEL_TYPE),
81                          ('time_index', time_index),]
82    if hostname:
83        equality_constraints.append(('hostname', hostname))
84    if labels:
85        for label in labels:
86            equality_constraints.append(('labels', label))
87    results = autotest_es.query(
88            fields_returned=['hostname', 'labels'],
89            equality_constraints=equality_constraints)
90
91    host_labels = {}
92    for hit in results.hits:
93        if 'labels' in hit:
94            host_labels[hit['hostname']] = hit['labels']
95
96    return host_labels
97
98
99def collect_info():
100    """Collect label info and report to metaDB.
101    """
102    # time_index is to index all host labels collected together. It's
103    # converted to int to make search faster.
104    time_index = int(time.time())
105    hosts = models.Host.objects.filter(invalid=False)
106    data_list = []
107    for host in hosts:
108        info = {'_type': _HOST_LABEL_TYPE,
109                'hostname': host.hostname,
110                'labels': [label.name for label in host.labels.all()],
111                'time_index': time_index}
112        data_list.append(info)
113    if not autotest_es.bulk_post(data_list, log_time_recorded=False):
114        raise Exception('Failed to upload host label info.')
115
116    # After all host label information is logged, save the time stamp.
117    autotest_es.post(use_http=True, type_str=_HOST_LABEL_TIME_INDEX_TYPE,
118                     metadata={'time_index': time_index},
119                     log_time_recorded=False)
120    logging.info('Finished collecting host labels for %d hosts.', len(hosts))
121
122
123def main():
124    """Main script.
125    """
126    parser = argparse.ArgumentParser()
127    parser.add_argument('action',
128                        help=('collect or query. Action collect will collect '
129                              'all hosts and their labels to metaDB. Action '
130                              'query will query for hosts and their labels '
131                              'information at a given day'))
132    parser.add_argument('-d', '--days_back', type=int, dest='days_back',
133                        help=('Number of days before current time. Query will '
134                              'get host label information collected before that'
135                              ' time. The option is applicable to query only. '
136                              'Default to 0, i.e., get the latest label info.'),
137                        default=0)
138    parser.add_argument('-n', '--hostname', type=str, dest='hostname',
139                        help=('Name of the host to query label information for.'
140                              'The option is applicable to query only. '
141                              'Default to None, i.e., return label info for all'
142                              ' hosts.'),
143                        default=None)
144    parser.add_argument('-l', '--labels', nargs='+', dest='labels',
145                        help=('A list of labels to filter hosts. The option is '
146                              'applicable to query only. Default to None.'),
147                        default=None)
148    parser.add_argument('-v', '--verbose', action="store_true", dest='verbose',
149                        help='Allow more detail information to be shown.')
150    options = parser.parse_args()
151
152    logging.getLogger().setLevel(logging.INFO if options.verbose
153                                 else logging.WARN)
154    if options.action == 'collect':
155        collect_info()
156    elif options.action == 'query':
157        host_labels = get_host_labels(options.days_back, options.hostname,
158                                      options.labels)
159        pprint.pprint(host_labels)
160    else:
161        logging.error('action %s is not supported, can only be collect or '
162                      'query!', options.action)
163
164
165if __name__ == '__main__':
166    main()
167