src/native_client/pynacl/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
  40 def HttpDownload(url, target):
  41   """Default download route."""
  42   http_download.HttpDownload(url, os.path.abspath(target), verbose=False,
  43                              logger=logging.debug)
  44
  45
  46 class GSDStorageError(Exception):
  47   """Error indicating writing to storage failed."""
  48   pass
  49
  50
  51 class GSDStorage(object):
  52   """A wrapper for reading and writing to GSD buckets.
  53
  54   Multiple read buckets may be specified, and the wrapper will sequentially try
  55   each and fall back to the next if the previous fails.
  56   Writing is to a single bucket.
  57   """
  58   def __init__(self,
  59                write_bucket,
  60                read_buckets,
  61                gsutil=None,
  62                call=subprocess.call,
  63                download=HttpDownload):
  64     """Init for this class.
  65
  66     Args:
  67       write_bucket: Google storage location to write to.
  68       read_buckets: Google storage locations to read from in preferred order.
  69       gsutil: List of cmd components needed to call gsutil.
  70       call: Testing hook to intercept command invocation.
  71       download: Testing hook to intercept download.
  72     """
  73     if gsutil is None:
  74       try:
  75         # Require that gsutil be Python if it is specified in the environment.
  76         gsutil = [sys.executable,
  77                   file_tools.Which(os.environ.get('GSUTIL', 'gsutil'),
  78                                    require_executable=False)]
  79       except file_tools.ExecutableNotFound:
  80         gsutil = ['gsutil']
  81     assert isinstance(gsutil, list)
  82     assert isinstance(read_buckets, list)
  83     self._gsutil = gsutil
  84     self._write_bucket = write_bucket
  85     self._read_buckets = read_buckets
  86     self._call = call
  87     self._download = download
  88
  89   def PutFile(self, path, key, clobber=True):
  90     """Write a file to global storage.
  91
  92     Args:
  93       path: Path of the file to write.
  94       key: Key to store file under.
  95     Raises:
  96       GSDStorageError if the underlying storage fails.
  97     Returns:
  98       URL written to.
  99     """
 100     if self._write_bucket is None:
 101       raise GSDStorageError('no bucket when storing %s to %s' % (path, key))
 102     obj = self._write_bucket + '/' + key
 103     arguments = ['cp', '-a', 'public-read']
 104     if not clobber:
 105       arguments.append('-n')
 106
 107     # Using file://c:/foo/bar form of path as gsutil does not like drive
 108     # letters without it.
 109     cmd = self._gsutil + arguments + [
 110         'file://' + os.path.abspath(path).replace(os.sep, '/'),
 111         GS_PATTERN % obj]
 112     logging.info('Running: %s' % str(cmd))
 113     if self._call(cmd) != 0:
 114       raise GSDStorageError('failed when storing %s to %s (%s)' % (
 115         path, key, cmd))
 116     return GS_HTTPS_PATTERN % obj
 117
 118   def PutData(self, data, key, clobber=True):
 119     """Write data to global storage.
 120
 121     Args:
 122       data: Data to store.
 123       key: Key to store file under.
 124     Raises:
 125       GSDStorageError if the underlying storage fails.
 126     Returns:
 127       URL written to.
 128     """
 129     handle, path = tempfile.mkstemp(prefix='gdstore', suffix='.tmp')
 130     try:
 131       os.close(handle)
 132       file_tools.WriteFile(data, path)
 133       return self.PutFile(path, key, clobber=clobber)
 134     finally:
 135       os.remove(path)
 136
 137   def GetFile(self, key, path):
 138     """Read a file from global storage.
 139
 140     Args:
 141       key: Key to store file under.
 142       path: Destination filename.
 143     Returns:
 144       URL used on success or None for failure.
 145     """
 146     for bucket in self._read_buckets:
 147       try:
 148         obj = bucket + '/' + key
 149         uri = GS_HTTPS_PATTERN % obj
 150         logging.debug('Downloading: %s to %s' % (uri, path))
 151         self._download(uri, path)
 152         return uri
 153       except:
 154         logging.debug('Failed downloading: %s to %s' % (uri, path))
 155     return None
 156
 157   def GetSecureFile(self, key, path):
 158     """ Read a non-publicly-accessible file from global storage.
 159
 160     Args:
 161       key: Key file is stored under.
 162       path: Destination filename
 163     Returns:
 164       command used on success or None on failure.
 165     """
 166     for bucket in self._read_buckets:
 167       try:
 168         obj = bucket + '/' + key
 169         cmd = self._gsutil + [
 170             'cp', GS_PATTERN % obj,
 171             'file://' + os.path.abspath(path).replace(os.sep, '/')]
 172         logging.info('Running: %s' % str(cmd))
 173         if self._call(cmd) == 0:
 174           return cmd
 175       except:
 176         logging.debug('Failed to fetch %s from %s (%s)' % (key, path, cmd))
 177     return None
 178
 179   def GetData(self, key):
 180     """Read data from global storage.
 181
 182     Args:
 183       key: Key to store file under.
 184     Returns:
 185       Data from storage, or None for failure.
 186     """
 187     work_dir = tempfile.mkdtemp(prefix='gdstore', suffix='.tmp')
 188     try:
 189       path = os.path.join(work_dir, 'data')
 190       if self.GetFile(key, path) is not None:
 191         return file_tools.ReadFile(path)
 192       return None
 193     finally:
 194       shutil.rmtree(work_dir)