• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
11from __future__ import print_function
12
13import logging
14import os
15
16from apiclient import discovery
17from apiclient import errors
18from oauth2client.client import ApplicationDefaultCredentialsError
19from oauth2client.client import GoogleCredentials
20
21# Cloud service
22PUBSUB_SERVICE_NAME = 'pubsub'
23PUBSUB_VERSION = 'v1beta2'
24PUBSUB_SCOPES = ['https://www.googleapis.com/auth/pubsub']
25# number of retry to publish an event.
26DEFAULT_PUBSUB_NUM_RETRIES = 3
27
28class PubSubException(Exception):
29    """Exception to be raised when the test to push to prod failed."""
30    pass
31
32
33class PubSubClient(object):
34    """A generic pubsub client."""
35    def __init__(self, credential_file=None):
36        """Constructor for PubSubClient.
37
38        Args:
39          credential_file: The credential filename.
40
41        Raises:
42          PubSubException if the credential file does not exist or corrupted.
43        """
44        if not credential_file:
45            raise PubSubException('You need to specify a credential file.')
46        self.credential_file = credential_file
47        self.credential = self._get_credential()
48
49    def _get_credential(self):
50        """Gets the pubsub service api handle."""
51        if not os.path.isfile(self.credential_file):
52            logging.error('No credential file found')
53            raise PubSubException('Credential file does not exist:' +
54                                  self.credential_file)
55        try:
56            credential = GoogleCredentials.from_stream(self.credential_file)
57            if credential.create_scoped_required():
58                credential = credential.create_scoped(PUBSUB_SCOPES)
59            return credential
60        except ApplicationDefaultCredentialsError as ex:
61            logging.exception('Failed to get credential:%s', ex)
62        except errors.Error as e:
63            logging.exception('Failed to get the pubsub service handle:%s', e)
64
65        raise PubSubException('Credential file %s does not exists:' %
66                              self.credential_file)
67
68    def _get_pubsub_service(self):
69        try:
70            return discovery.build(PUBSUB_SERVICE_NAME, PUBSUB_VERSION,
71                                   credentials=self.credential)
72        except errors.Error as e:
73            logging.exception('Failed to get pubsub resource object:%s', e)
74            raise PubSubException('Failed to get pubsub resource object')
75
76    def publish_notifications(self, topic, messages=None):
77        """Publishes a test result notification to a given pubsub topic.
78
79        @param topic: The Cloud pubsub topic.
80        @param messages: A list of notification messages.
81
82        @returns A list of pubsub message ids, and empty if fails.
83
84        @raises PubSubException if failed to publish the notification.
85        """
86        if not messages:
87            return None
88
89        pubsub = self._get_pubsub_service()
90        try:
91            body = {'messages': messages}
92            resp = pubsub.projects().topics().publish(
93                topic=topic, body=body).execute(
94                    num_retries=DEFAULT_PUBSUB_NUM_RETRIES)
95            msgIds = []
96            if resp:
97                msgIds = resp.get('messageIds')
98                if msgIds:
99                    logging.debug('Published notification message')
100                else:
101                    logging.error('Failed to published notification message')
102            return msgIds
103        except errors.Error as e:
104            logging.exception('Failed to publish test result notification:%s',
105                    e)
106            raise PubSubException('Failed to publish the notification')
107