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