site_utils/add_detected_host_labels.py

#!/usr/bin/python

# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""A small wrapper script, iterates through
the known hosts and tries to call get_labels()
to discover host functionality, and adds these
detected labels to host.

Limitations:
 - Does not keep a count of how many labels were
   actually added.
 - If a label is added by this script because it
   is detected as supported by get_labels, but later becomes
   unsupported, this script has no way to know that it
   should be removed, so it will remain attached to the host.
   See crosbug.com/38569
"""


from multiprocessing import pool
import logging
import socket
import argparse
import sys

import common

from autotest_lib.server import hosts
from autotest_lib.server import frontend
from autotest_lib.client.common_lib import error


def add_missing_labels(afe, hostname):
    """
    Queries the detectable labels supported by the given host,
    and adds those labels to the host.

    @param afe: A frontend.AFE() instance.
    @param hostname: The host to query and update.

    @return: True on success.
             False on failure to fetch labels or to add any individual label.
    """

    try:
        host = hosts.create_host(hostname)
        labels = host.get_labels()
    except socket.gaierror:
        logging.warning('Unable to establish ssh connection to hostname '
                        '%s. Skipping.', hostname)
        return False
    except error.AutoservError:
        logging.warning('Unable to query labels on hostname %s. Skipping.',
                         hostname)
        return False
    finally:
        host.close()

    label_matches = afe.get_labels(name__in=labels)

    for label in label_matches:
        label.add_hosts(hosts=[hostname])

    missing_labels = set(labels) - set([l.name for l in label_matches])

    if missing_labels:
        for label in missing_labels:
            logging.warning('Unable to add label %s to host %s. '
                            'Skipping unknown label.', label, hostname)
        return False

    return True


def main():
    """"
    Entry point for add_detected_host_labels script.
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--silent', dest='silent', action='store_true',
                        help='Suppress all but critical logging messages.')
    parser.add_argument('-i', '--info', dest='info_only', action='store_true',
                        help='Suppress logging messages below INFO priority.')
    parser.add_argument('-m', '--machines', dest='machines',
                        help='Comma separated list of machines to check.')
    options = parser.parse_args()

    if options.silent and options.info_only:
        print 'The -i and -s flags cannot be used together.'
        parser.print_help()
        return 0


    if options.silent:
        logging.disable(logging.CRITICAL)

    if options.info_only:
        logging.disable(logging.DEBUG)

    threadpool = pool.ThreadPool()
    afe = frontend.AFE()

    if options.machines:
        hostnames = [m.strip() for m in options.machines.split(',')]
    else:
        hostnames = afe.get_hostnames()

    successes = sum(threadpool.imap_unordered(
                        lambda x: add_missing_labels(afe, x),
                        hostnames))
    attempts = len(hostnames)

    logging.info('Label updating finished. Failed update on %d out of %d '
                 'hosts.', attempts-successes, attempts)

    return 0


if __name__ == '__main__':
    sys.exit(main())