Home | History | Annotate | Download | only in site_utils 
      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 
      7 Uses gsutil to archive files to the configured Google Storage bucket.
      8 Upon successful copy, the local results directory is deleted.
      9 """
     10 
     11 from __future__ import print_function
     12 
     13 import os
     14 
     15 from apiclient import discovery
     16 from apiclient import errors
     17 from oauth2client.client import ApplicationDefaultCredentialsError
     18 from oauth2client.client import GoogleCredentials
     19 from chromite.lib import cros_logging as logging
     20 
     21 # Cloud service
     22 PUBSUB_SERVICE_NAME = 'pubsub'
     23 PUBSUB_VERSION = 'v1beta2'
     24 PUBSUB_SCOPES = ['https://www.googleapis.com/auth/pubsub']
     25 # number of retry to publish an event.
     26 DEFAULT_PUBSUB_NUM_RETRIES = 3
     27 
     28 class PubSubException(Exception):
     29     """Exception to be raised when the test to push to prod failed."""
     30     pass
     31 
     32 
     33 class 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