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 """ 8 This script provides functions to: 9 1. 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 12 2. 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 16 import argparse 17 import itertools 18 import logging 19 import pprint 20 import time 21 22 import common 23 from autotest_lib.client.common_lib import time_utils 24 from autotest_lib.client.common_lib.cros.graphite import autotest_es 25 from autotest_lib.frontend import setup_django_environment 26 from 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 33 def 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 50 def 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 99 def 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 123 def 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 165 if __name__ == '__main__': 166 main() 167
