tools/android/loading Split GoogleAPI wrappers to a separate library

tools/android/loading/gce/main.py

 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import json
 import os
 import re
 import threading
 import time
 import subprocess
 import sys
 
-from gcloud import storage
-from gcloud.exceptions import NotFound
-from oauth2client.client import GoogleCredentials
-
 # NOTE: The parent directory needs to be first in sys.path to avoid conflicts
 # with catapult modules that have colliding names, as catapult inserts itself
 # into the path as the second element. This is an ugly and fragile hack.
 sys.path.insert(0,
     os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
 import controller
+from google_api_util import GoogleAPIUtil
 import loading_trace
+from loading_trace_database import LoadingTraceDatabase
 import options
-from loading_trace_database import LoadingTraceDatabase
 
 
 class ServerApp(object):
   """Simple web server application, collecting traces and writing them in
   Google Cloud Storage.
   """
 
   def __init__(self, configuration_file):
     """|configuration_file| is a path to a file containing JSON as described in
     README.md.
     """
     self._tasks = []  # List of remaining tasks, only modified by _thread.
     self._failed_tasks = []  # Failed tasks, only modified by _thread.
     self._thread = None
     self._tasks_lock = threading.Lock()  # Protects _tasks and _failed_tasks.
     self._initial_task_count = -1
     self._start_time = None
-    print 'Initializing credentials'
-    self._credentials = GoogleCredentials.get_application_default()
     print 'Reading configuration'
     with open(configuration_file) as config_json:
        config = json.load(config_json)
-       self._project_name = config['project_name']
 
        # Separate the cloud storage path into the bucket and the base path under
        # the bucket.
        storage_path_components = config['cloud_storage_path'].split('/')
        self._bucket_name = storage_path_components[0]
        self._base_path_in_bucket = ''
        if len(storage_path_components) > 1:
          self._base_path_in_bucket = '/'.join(storage_path_components[1:])
          if not self._base_path_in_bucket.endswith('/'):
            self._base_path_in_bucket += '/'
 
        self._src_path = config['src_path']
+       self._google_api_util = GoogleAPIUtil(
+           project_name=config['project_name'], bucket_name=self._bucket_name)
 
     # Initialize the global options that will be used during trace generation.
     options.OPTIONS.ParseArgs([])
     options.OPTIONS.local_binary = config['chrome_path']
 
   def _IsProcessingTasks(self):
     """Returns True if the application is currently processing tasks."""
     return self._thread is not None and self._thread.is_alive()
 
-  def _GetStorageClient(self):
-    return storage.Client(project = self._project_name,
-                          credentials = self._credentials)
-
-  def _GetStorageBucket(self, storage_client):
-    return storage_client.get_bucket(self._bucket_name)
-
-  def _UploadFile(self, filename_src, filename_dest):
-    """Uploads a file to Google Cloud Storage
-
-    Args:
-      filename_src: name of the local file
-      filename_dest: name of the file in Google Cloud Storage
-
-    Returns:
-      The URL of the file in Google Cloud Storage.
-    """
-    client = self._GetStorageClient()
-    bucket = self._GetStorageBucket(client)
-    blob = bucket.blob(filename_dest)
-    with open(filename_src) as file_src:
-      blob.upload_from_file(file_src)
-    return blob.public_url
-
-  def _UploadString(self, data_string, filename_dest):
-    """Uploads a string to Google Cloud Storage
-
-    Args:
-      data_string: the contents of the file to be uploaded
-      filename_dest: name of the file in Google Cloud Storage
-
-    Returns:
-      The URL of the file in Google Cloud Storage.
-    """
-    client = self._GetStorageClient()
-    bucket = self._GetStorageBucket(client)
-    blob = bucket.blob(filename_dest)
-    blob.upload_from_string(data_string)
-    return blob.public_url
-
   def _GenerateTrace(self, url, emulate_device, emulate_network, filename,
                      log_filename):
     """ Generates a trace on _thread.
 
     Args:
       url: URL as a string.
       emulate_device: Name of the device to emulate. Empty for no emulation.
       emulate_network: Type of network emulation. Empty for no emulation.
       filename: Name of the file where the trace is saved.
       log_filename: Name of the file where standard output and errors are logged
@@ skipping 83 matching lines @@
       url = self._tasks[-1]
       local_filename = pattern.sub('_', url)
       for repeat in range(repeat_count):
         print 'Generating trace for URL: %s' % url
         remote_filename = local_filename + '/' + str(repeat)
         trace_metadata = self._GenerateTrace(
             url, emulate_device, emulate_network, local_filename, log_filename)
         if trace_metadata['succeeded']:
           print 'Uploading: %s' % remote_filename
           remote_trace_location = traces_dir + remote_filename
-          self._UploadFile(local_filename, remote_trace_location)
+          self._google_api_util.UploadFile(local_filename,
+                                           remote_trace_location)
           full_cloud_storage_path = ('gs://' + self._bucket_name + '/' +
               remote_trace_location)
           trace_database.AddTrace(full_cloud_storage_path, trace_metadata)
         else:
           print 'Trace generation failed for URL: %s' % url
           self._tasks_lock.acquire()
           self._failed_tasks.append({ "url": url, "repeat": repeat})
           self._tasks_lock.release()
           if os.path.isfile(local_filename):
-            self._UploadFile(local_filename, failures_dir + remote_filename)
+            self._google_api_util.UploadFile(local_filename,
+                                            failures_dir + remote_filename)
         print 'Uploading log'
-        self._UploadFile(log_filename, logs_dir + remote_filename)
+        self._google_api_util.UploadFile(log_filename,
+                                         logs_dir + remote_filename)
       # Pop once task is finished, for accurate status tracking.
       self._tasks_lock.acquire()
       url = self._tasks.pop()
       self._tasks_lock.release()
 
-    self._UploadString(json.dumps(trace_database.ToJsonDict(), indent=2),
-                       traces_dir + 'trace_database.json')
+    self._google_api_util.UploadString(
+        json.dumps(trace_database.ToJsonDict(), indent=2),
+        traces_dir + 'trace_database.json')
 
     if len(self._failed_tasks) > 0:
       print 'Uploading failing URLs'
-      self._UploadString(json.dumps(self._failed_tasks, indent=2),
-                         failures_dir + 'failures.json')
+      self._google_api_util.UploadString(
+          json.dumps(self._failed_tasks, indent=2),
+          failures_dir + 'failures.json')
 
   def _SetTaskList(self, http_body):
     """Sets the list of tasks and starts processing them
 
     Args:
       http_body: JSON dictionary. See README.md for a description of the format.
 
     Returns:
       A string to be sent back to the client, describing the success status of
       the request.
@@ skipping 62 matching lines @@
     response_headers = [
         ('Content-type','text/plain'),
         ('Content-Length', str(len(data)))
     ]
     start_response('200 OK', response_headers)
     return iter([data])
 
 
 def StartApp(configuration_file):
   return ServerApp(configuration_file)