gs_utils.upload_file(): upload with temp name, verify contents, then rename

py/utils/gs_utils.py

 #!/usr/bin/python
 
 # pylint: disable=C0301
 """
 Copyright 2014 Google Inc.
 
 Use of this source code is governed by a BSD-style license that can be
 found in the LICENSE file.
 
 Utilities for accessing Google Cloud Storage, using the boto library (wrapper
@@ skipping 142 matching lines @@
       bucket: GS bucket to delete a file from
       path: full path (Posix-style) of the file within the bucket to delete
     """
     b = self._connect_to_bucket(bucket=bucket)
     key = Key(b)
     key.name = path
     try:
       key.delete()
     except BotoServerError, e:
       e.body = (repr(e.body) +
-                ' while deleting bucket=%s, path=%s' % (b.name, path))
+                ' while deleting gs://%s/%s' % (b.name, path))
       raise
 
   def get_last_modified_time(self, bucket, path):
     """Gets the timestamp of when this file was last modified.
 
     Params:
       bucket: GS bucket in which to look for the file
       path: full path (Posix-style) of the file within the bucket to check
 
     Returns the last modified time, as a freeform string.  If the file was not
     found, returns None.
     """
     b = self._connect_to_bucket(bucket=bucket)
     try:
       key = b.get_key(key_name=path)
       if not key:
         return None
       return key.last_modified
     except BotoServerError, e:
       e.body = (repr(e.body) +
-                ' while getting attributes of bucket=%s, path=%s' % (
-                    b.name, path))
+                ' while getting attributes of gs://%s/%s' % (b.name, path))
       raise
 
   def upload_file(self, source_path, dest_bucket, dest_path,
                   upload_if=UploadIf.ALWAYS,
                   predefined_acl=None,
                   fine_grained_acl_list=None):
     """Upload contents of a local file to Google Storage.
 
     params:
       source_path: full path (local-OS-style) on local disk to read from
       dest_bucket: GS bucket to copy the file to
       dest_path: full path (Posix-style) within that bucket
       upload_if: one of the UploadIf values, describing in which cases we should
           upload the file
       predefined_acl: which predefined ACL to apply to the file on Google
           Storage; must be one of the PredefinedACL values defined above.
           If None, inherits dest_bucket's default object ACL.
       fine_grained_acl_list: list of (id_type, id_value, permission) tuples
           to apply to the uploaded file (on top of the predefined_acl),
           or None if predefined_acl is sufficient
 
     TODO(epoger): Consider adding a do_compress parameter that would compress
     the file using gzip before upload, and add a "Content-Encoding:gzip" header
     so that HTTP downloads of the file would be unzipped automatically.
     See https://developers.google.com/storage/docs/gsutil/addlhelp/
         WorkingWithObjectMetadata#content-encoding
     """
     b = self._connect_to_bucket(bucket=dest_bucket)
+    local_md5 = None  # filled in lazily
 
     if upload_if == self.UploadIf.IF_NEW:
       old_key = b.get_key(key_name=dest_path)
       if old_key:
         print 'Skipping upload of existing file gs://%s/%s' % (
             b.name, dest_path)
         return
     elif upload_if == self.UploadIf.IF_MODIFIED:
       old_key = b.get_key(key_name=dest_path)
       if old_key:
-        local_md5 = '"%s"' % _get_local_md5(path=source_path)
-        if local_md5 == old_key.etag:
+        if not local_md5:
+          local_md5 = _get_local_md5(path=source_path)
+        if ('"%s"' % local_md5) == old_key.etag:
           print 'Skipping upload of unmodified file gs://%s/%s : %s' % (
               b.name, dest_path, local_md5)
           return
     elif upload_if != self.UploadIf.ALWAYS:
       raise Exception('unknown value of upload_if: %s' % upload_if)
 
-    key = Key(b)
-    key.name = dest_path
+    # Upload the file using a temporary name at first, in case the transfer
+    # is interrupted partway through.
+    if not local_md5:
+      local_md5 = _get_local_md5(path=source_path)
+    initial_key = Key(b)
+    initial_key.name = dest_path + '-uploading-' + local_md5
     try:
-      key.set_contents_from_filename(filename=source_path,
-                                     policy=predefined_acl)
+      initial_key.set_contents_from_filename(filename=source_path,
+                                             policy=predefined_acl)
     except BotoServerError, e:
       e.body = (repr(e.body) +
-                ' while uploading source_path=%s to bucket=%s, path=%s' % (
-                    source_path, b.name, key.name))
+                ' while uploading source_path=%s to gs://%s/%s' % (
+                    source_path, b.name, initial_key.name))
       raise
+
+    # Verify that the file contents were uploaded successfully.
+    #
+    # TODO(epoger): Check whether the boto library or XML API already do this...
+    # if so, we may be duplicating effort here, and maybe we don't need to do
+    # the whole "upload using temporary filename, then rename" thing.
+    #
+    # TODO(epoger): Confirm that the etag is set on the server side...
+    # otherwise, we may just be validating another MD5 hash that was generated
+    # on the client side before the file was uploaded!
+    validate_key = b.get_key(key_name=initial_key.name)
+    if validate_key.etag != ('"%s"' % local_md5):
+      raise Exception('found wrong MD5 after uploading gs://%s/%s' % (
+          b.name, validate_key.name))
+
+    # Rename the file to its real name.
+    #
+    # TODO(epoger): I don't know how long this takes.  I wish we could rename
+    # the key instead, but AFAICT you can't do that.
+    # Perhaps we could use Key.compose() to create a composite object pointing
+    # at the original key?
+    # See https://developers.google.com/storage/docs/composite-objects
+    final_key = b.copy_key(
+        new_key_name=dest_path, src_key_name=initial_key.name,
+        src_bucket_name=b.name, preserve_acl=False)
+    initial_key.delete()
+
+    # Set ACLs on the file.
+    # We do this *after* copy_key(), because copy_key's preserve_acl
+    # functionality would incur a performance hit.
     for (id_type, id_value, permission) in fine_grained_acl_list or []:
       self.set_acl(
-          bucket=b, path=key.name,
+          bucket=b, path=final_key.name,
           id_type=id_type, id_value=id_value, permission=permission)
 
   def upload_dir_contents(self, source_dir, dest_bucket, dest_dir,
                           upload_if=UploadIf.ALWAYS, **kwargs):
     """Recursively upload contents of a local directory to Google Storage.
 
     params:
       source_dir: full path (local-OS-style) on local disk of directory to copy
           contents of
       dest_bucket: GS bucket to copy the files into
@@ skipping 75 matching lines @@
     b = self._connect_to_bucket(bucket=source_bucket)
     key = Key(b)
     key.name = source_path
     if create_subdirs_if_needed:
       _makedirs_if_needed(os.path.dirname(dest_path))
     with open(dest_path, 'w') as f:
       try:
         key.get_contents_to_file(fp=f)
       except BotoServerError, e:
         e.body = (repr(e.body) +
-                  ' while downloading bucket=%s, path=%s to local_path=%s' % (
+                  ' while downloading gs://%s/%s to local_path=%s' % (
                       b.name, source_path, dest_path))
         raise
 
   def download_dir_contents(self, source_bucket, source_dir, dest_dir):
     """Recursively download contents of a Google Storage directory to local disk
 
     params:
       source_bucket: GS bucket to copy the files from
       source_dir: full path (Posix-style) within that bucket; read the files
           from this directory
@@ skipping 13 matching lines @@
 
     for filename in files:
       key = Key(b)
       key.name = posixpath.join(source_dir, filename)
       dest_path = os.path.join(dest_dir, filename)
       with open(dest_path, 'w') as f:
         try:
           key.get_contents_to_file(fp=f)
         except BotoServerError, e:
           e.body = (repr(e.body) +
-                    ' while downloading bucket=%s, path=%s to local_path=%s' % (
+                    ' while downloading gs://%s/%s to local_path=%s' % (
                         b.name, key.name, dest_path))
           raise
 
     for dirname in dirs:
       self.download_dir_contents(  # recurse
           source_bucket=source_bucket,
           source_dir=posixpath.join(source_dir, dirname),
           dest_dir=os.path.join(dest_dir, dirname))
 
   def get_acl(self, bucket, path, id_type, id_value):
@@ skipping 183 matching lines @@
 
 def _get_local_md5(path):
   """Returns the MD5 hash of a file on local disk."""
   hasher = hashlib.md5()
   with open(path, 'rb') as f:
     while True:
       data = f.read(64*1024)
       if not data:
         return hasher.hexdigest()
       hasher.update(data)