1# Copyright 2016 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"""Script to archive old Autotest results to Google Storage.
6
7Uses gsutil to archive files to the configured Google Storage bucket.
8Upon successful copy, the local results directory is deleted.
9"""
10
11import logging
12import os
13
14from apiclient import discovery
15from oauth2client.client import ApplicationDefaultCredentialsError
16from oauth2client.client import GoogleCredentials
17from autotest_lib.server.hosts import moblab_host
18
19import common
20
21# Cloud service
22# TODO(ntang): move this to config.
23CLOUD_SERVICE_ACCOUNT_FILE = moblab_host.MOBLAB_SERVICE_ACCOUNT_LOCATION
24PUBSUB_SERVICE_NAME = 'pubsub'
25PUBSUB_VERSION = 'v1beta2'
26PUBSUB_SCOPES = ['https://www.googleapis.com/auth/pubsub']
27# number of retry to publish an event.
28_PUBSUB_NUM_RETRIES = 3
29
30
31def _get_pubsub_service():
32    """Gets the pubsub service api handle."""
33    if not os.path.isfile(CLOUD_SERVICE_ACCOUNT_FILE):
34        logging.error('No credential file found')
35        return None
36
37    try:
38        credentials = GoogleCredentials.from_stream(CLOUD_SERVICE_ACCOUNT_FILE)
39        if credentials.create_scoped_required():
40            credentials = credentials.create_scoped(PUBSUB_SCOPES)
41        return discovery.build(PUBSUB_SERVICE_NAME, PUBSUB_VERSION,
42                                 credentials=credentials)
43    except ApplicationDefaultCredentialsError as ex:
44        logging.error('Failed to get credential.')
45    except:
46        logging.error('Failed to get the pubsub service handle.')
47
48    return None
49
50
51def publish_notifications(topic, messages=[]):
52    """Publishes a test result notification to a given pubsub topic.
53
54    @param topic: The Cloud pubsub topic.
55    @param messages: A list of notification messages.
56
57    @returns A list of pubsub message ids, and empty if fails.
58    """
59    pubsub = _get_pubsub_service()
60    if pubsub:
61        try:
62            body = {'messages': messages}
63            resp = pubsub.projects().topics().publish(topic=topic,
64                    body=body).execute(num_retries=_PUBSUB_NUM_RETRIES)
65            if resp:
66                msgIds = resp.get('messageIds')
67                if msgIds:
68                    logging.debug('Published notification message')
69                    return msgIds
70        except:
71            pass
72    logging.error('Failed to publish test result notifiation.')
73    return []
74