src/native_client/build/gsd_storage.py

   1 #!/usr/bin/python
   2 # Copyright (c) 2012 The Native Client Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """Provide Google Storage access.
   7
   8 Provide an high-level interface to Google Storage.
   9 Operations are provided to read/write whole files and to
  10 read/write strings. This allows google storage to be treated
  11 more or less like a key+value data-store.
  12 """
  13
  14
  15 import logging
  16 import os
  17 import re
  18 import shutil
  19 import subprocess
  20 import sys
  21 import tempfile
  22
  23 import file_tools
  24 import http_download
  25
  26
  27 GS_PATTERN = 'gs://%s'
  28 GS_HTTPS_PATTERN = 'https://storage.googleapis.com/%s'
  29
  30
  31 def LegalizeName(name):
  32   """ Return a file name suitable for uploading to Google Storage.
  33
  34   The names of such files cannot contain dashes or other non-identifier
  35   characters.
  36   """
  37   return re.sub(r'[^A-Za-z0-9_/.]', '_', name)
  38
  39 def HttpDownload(url, target):
  40   """Default download route."""
  41   http_download.HttpDownload(url, os.path.abspath(target),
  42                              logger=logging.info)
  43
  44
  45 class GSDStorageError(Exception):
  46   """Error indicating writing to storage failed."""
  47   pass
  48
  49
  50 class GSDStorage(object):
  51   """A wrapper for reading and writing to GSD buckets.
  52
  53   Multiple read buckets may be specified, and the wrapper will sequentially try
  54   each and fall back to the next if the previous fails.
  55   Writing is to a single bucket.
  56   """
  57   def __init__(self,
  58                write_bucket,
  59                read_buckets,
  60                gsutil=None,
  61                call=subprocess.call,
  62                download=HttpDownload):
  63     """Init for this class.
  64
  65     Args:
  66       write_bucket: Google storage location to write to.
  67       read_buckets: Google storage locations to read from in preferred order.
  68       gsutil: List of cmd components needed to call gsutil.
  69       call: Testing hook to intercept command invocation.
  70       download: Testing hook to intercept download.
  71     """
  72     if gsutil is None:
  73       try:
  74         # Require that gsutil be Python if it is specified in the environment.
  75         gsutil = [sys.executable,
  76                   file_tools.Which(os.environ.get('GSUTIL', 'gsutil'),
  77                                    require_executable=False)]
  78       except file_tools.ExecutableNotFound:
  79         gsutil = ['gsutil']
  80     assert isinstance(gsutil, list)
  81     assert isinstance(read_buckets, list)
  82     self._gsutil = gsutil
  83     self._write_bucket = write_bucket
  84     self._read_buckets = read_buckets
  85     self._call = call
  86     self._download = download
  87
  88   def PutFile(self, path, key):
  89     """Write a file to global storage.
  90
  91     Args:
  92       path: Path of the file to write.
  93       key: Key to store file under.
  94     Raises:
  95       GSDStorageError if the underlying storage fails.
  96     Returns:
  97       URL written to.
  98     """
  99     if self._write_bucket is None:
 100       raise GSDStorageError('no bucket when storing %s to %s' % (path, key))
 101     obj = self._write_bucket + '/' + key
 102     # Using file://c:/foo/bar form of path as gsutil does not like drive
 103     # letters without it.
 104     cmd = self._gsutil + [
 105         'cp', '-a', 'public-read',
 106         'file://' + os.path.abspath(path).replace(os.sep, '/'),
 107         GS_PATTERN % obj]
 108     logging.info('Running: %s' % str(cmd))
 109     if self._call(cmd) != 0:
 110       raise GSDStorageError('failed when storing %s to %s (%s)' % (
 111         path, key, cmd))
 112     return GS_HTTPS_PATTERN % obj
 113
 114   def PutData(self, data, key):
 115     """Write data to global storage.
 116
 117     Args:
 118       data: Data to store.
 119       key: Key to store file under.
 120     Raises:
 121       GSDStorageError if the underlying storage fails.
 122     Returns:
 123       URL written to.
 124     """
 125     handle, path = tempfile.mkstemp(prefix='gdstore', suffix='.tmp')
 126     try:
 127       os.close(handle)
 128       file_tools.WriteFile(data, path)
 129       return self.PutFile(path, key)
 130     finally:
 131       os.remove(path)
 132
 133   def GetFile(self, key, path):
 134     """Read a file from global storage.
 135
 136     Args:
 137       key: Key to store file under.
 138       path: Destination filename.
 139     Returns:
 140       URL used on success or None for failure.
 141     """
 142     for bucket in self._read_buckets:
 143       try:
 144         obj = bucket + '/' + key
 145         uri = GS_HTTPS_PATTERN % obj
 146         logging.debug('Downloading: %s to %s' % (uri, path))
 147         self._download(uri, path)
 148         return uri
 149       except:
 150         logging.debug('Failed downloading: %s to %s' % (uri, path))
 151     return None
 152
 153   def GetSecureFile(self, key, path):
 154     """ Read a non-publicly-accessible file from global storage.
 155
 156     Args:
 157       key: Key file is stored under.
 158       path: Destination filename
 159     Returns:
 160       command used on success or None on failure.
 161     """
 162     for bucket in self._read_buckets:
 163       try:
 164         obj = bucket + '/' + key
 165         cmd = self._gsutil + [
 166             'cp', GS_PATTERN % obj,
 167             'file://' + os.path.abspath(path).replace(os.sep, '/')]
 168         logging.info('Running: %s' % str(cmd))
 169         if self._call(cmd) == 0:
 170           return cmd
 171       except:
 172         logging.debug('Failed to fetch %s from %s (%s)' % (key, path, cmd))
 173     return None
 174
 175   def GetData(self, key):
 176     """Read data from global storage.
 177
 178     Args:
 179       key: Key to store file under.
 180     Returns:
 181       Data from storage, or None for failure.
 182     """
 183     work_dir = tempfile.mkdtemp(prefix='gdstore', suffix='.tmp')
 184     try:
 185       path = os.path.join(work_dir, 'data')
 186       if self.GetFile(key, path) is not None:
 187         return file_tools.ReadFile(path)
 188       return None
 189     finally:
 190       shutil.rmtree(work_dir)