gs_utils: when uploading a whole dir using IF_NEW, check for existence of multiple files in a single operation BUG=skia:2778 R=borenet@google.com Review URL: https://codereview.chromium.org/420553002
diff --git a/py/utils/gs_utils.py b/py/utils/gs_utils.py index c58bc67..7f3b59c 100644 --- a/py/utils/gs_utils.py +++ b/py/utils/gs_utils.py
@@ -104,9 +104,9 @@ class UploadIf: """Cases in which we will upload a file. - Beware of performance tradeoffs. E.g., if the file is small, the extra - round trip to check for file existence and/or checksum may take longer than - just uploading the file. + Beware of performance tradeoffs. E.g., if you are uploading just one small + file, the extra round trip to check for file existence and/or checksum may + take longer than just uploading the file. See http://skbug.com/2778 ('gs_utils: when uploading IF_NEW, batch up checks for existing files within a single remote directory') """ @@ -244,7 +244,8 @@ bucket=b, path=key.name, id_type=id_type, id_value=id_value, permission=permission) - def upload_dir_contents(self, source_dir, dest_bucket, dest_dir, **kwargs): + 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: @@ -253,34 +254,67 @@ dest_bucket: GS bucket to copy the files into dest_dir: full path (Posix-style) within that bucket; write the files into this directory. If None, write into the root directory of the bucket. + upload_if: one of the UploadIf values, describing in which cases we should + upload the file kwargs: any additional keyword arguments "inherited" from upload_file() The copy operates as a merge: any files in source_dir will be "overlaid" on top of the existing content in dest_dir. Existing files with the same names - may or may not be overwritten, depending on the value of the upload_if kwarg - inherited from upload_file(). + may or may not be overwritten, depending on the value of upload_if. TODO(epoger): Upload multiple files simultaneously to reduce latency. - - TODO(epoger): When upload_if==IF_NEW, batch up checks for existing files - within a single remote directory. See http://skbug.com/2778 """ b = self._connect_to_bucket(bucket=dest_bucket) - for filename in sorted(os.listdir(source_dir)): - local_path = os.path.join(source_dir, filename) - if dest_dir: - remote_path = posixpath.join(dest_dir, filename) - else: - remote_path = filename + if not dest_dir: + dest_dir = '' - if os.path.isdir(local_path): - self.upload_dir_contents( # recurse - source_dir=local_path, dest_bucket=b, dest_dir=remote_path, - **kwargs) + # Create a set of all files within source_dir. + source_fileset = set() + prefix_length = len(source_dir)+1 + for dirpath, _, filenames in os.walk(source_dir): + relative_dirpath = dirpath[prefix_length:] + for filename in filenames: + source_fileset.add(os.path.join(relative_dirpath, filename)) + + # If we are only uploading files conditionally, remove any unnecessary + # files from source_fileset. + if upload_if == self.UploadIf.ALWAYS: + pass # there are no shortcuts... upload them all + else: + # Create a mapping of filename to Key for existing files within dest_dir + existing_dest_filemap = {} + prefix = dest_dir + if prefix and not prefix.endswith('/'): + prefix += '/' + prefix_length = len(prefix) + items = BucketListResultSet(bucket=b, prefix=prefix) + for item in items: + if type(item) is Key: + existing_dest_filemap[item.name[prefix_length:]] = item + + # Now, depending on upload_if, trim files we should skip uploading. + files_in_common = source_fileset.intersection( + existing_dest_filemap.keys()) + if upload_if == self.UploadIf.IF_NEW: + source_fileset -= files_in_common + elif upload_if == self.UploadIf.IF_MODIFIED: + for rel_path in files_in_common: + local_md5 = '"%s"' % _get_local_md5(path=os.path.join( + source_dir, rel_path)) + key = existing_dest_filemap[rel_path] + if local_md5 == key.etag: + source_fileset.remove(rel_path) else: - self.upload_file( - source_path=local_path, dest_bucket=b, dest_path=remote_path, - **kwargs) + raise Exception('unknown value of upload_if: %s' % upload_if) + + # Upload any files still in source_fileset. + for rel_path in sorted(source_fileset): + self.upload_file( + source_path=os.path.join(source_dir, rel_path), + dest_bucket=b, + dest_path=posixpath.join(dest_dir, rel_path), + upload_if=self.UploadIf.ALWAYS, + **kwargs) def download_file(self, source_bucket, source_path, dest_path, create_subdirs_if_needed=False):
diff --git a/py/utils/gs_utils_manualtest.py b/py/utils/gs_utils_manualtest.py index b4f70a6..37621d1 100755 --- a/py/utils/gs_utils_manualtest.py +++ b/py/utils/gs_utils_manualtest.py
@@ -45,7 +45,7 @@ gs.list_bucket_contents(bucket=TEST_BUCKET, subdir=None) -def _test_upload_if(): +def _test_upload_if_one_file(): """Test upload_if param within upload_file().""" gs = _get_authenticated_gs_handle() filename = 'filename' @@ -113,6 +113,86 @@ # Clean up the local dir. shutil.rmtree(local_dir) + +def _test_upload_if_multiple_files(): + """Test upload_if param within upload_dir_contents().""" + gs = _get_authenticated_gs_handle() + subdir = 'subdir' + filenames = ['file1', 'file2'] + local_dir = tempfile.mkdtemp() + remote_dir = _get_unique_posix_dir() + sample_file_local_path = os.path.join(local_dir, subdir, filenames[0]) + sample_file_remote_path = posixpath.join(remote_dir, subdir, filenames[0]) + try: + # Create files on local disk, and upload them for the first time. + os.mkdir(os.path.join(local_dir, subdir)) + for filename in filenames: + with open(os.path.join(local_dir, subdir, filename), 'w') as f: + f.write('original contents of %s' % filename) + gs.upload_dir_contents( + source_dir=local_dir, dest_bucket=TEST_BUCKET, + dest_dir=remote_dir, upload_if=gs.UploadIf.IF_NEW) + try: + # Re-upload the same files, with upload_if=gs.UploadIf.ALWAYS; + # the timestamps should change. + old_timestamp = gs.get_last_modified_time( + bucket=TEST_BUCKET, path=sample_file_remote_path) + time.sleep(2) + gs.upload_dir_contents( + source_dir=local_dir, dest_bucket=TEST_BUCKET, + dest_dir=remote_dir, upload_if=gs.UploadIf.ALWAYS) + new_timestamp = gs.get_last_modified_time( + bucket=TEST_BUCKET, path=sample_file_remote_path) + assert old_timestamp != new_timestamp, '%s != %s' % ( + old_timestamp, new_timestamp) + + # Re-upload the same files, with upload_if=gs.UploadIf.IF_MODIFIED; + # the timestamps should NOT change. + old_timestamp = new_timestamp + time.sleep(2) + gs.upload_dir_contents( + source_dir=local_dir, dest_bucket=TEST_BUCKET, + dest_dir=remote_dir, upload_if=gs.UploadIf.IF_MODIFIED) + new_timestamp = gs.get_last_modified_time( + bucket=TEST_BUCKET, path=sample_file_remote_path) + assert old_timestamp == new_timestamp, '%s == %s' % ( + old_timestamp, new_timestamp) + + # Modify and re-upload the files, with upload_if=gs.UploadIf.IF_NEW; + # the timestamps should still not change. + old_timestamp = new_timestamp + with open(sample_file_local_path, 'w') as f: + f.write('modified contents of sample file') + time.sleep(2) + gs.upload_dir_contents( + source_dir=local_dir, dest_bucket=TEST_BUCKET, + dest_dir=remote_dir, upload_if=gs.UploadIf.IF_NEW) + new_timestamp = gs.get_last_modified_time( + bucket=TEST_BUCKET, path=sample_file_remote_path) + assert old_timestamp == new_timestamp, '%s == %s' % ( + old_timestamp, new_timestamp) + + # Re-upload the modified file, with upload_if=gs.UploadIf.IF_MODIFIED; + # now the timestamp SHOULD change. + old_timestamp = new_timestamp + time.sleep(2) + gs.upload_dir_contents( + source_dir=local_dir, dest_bucket=TEST_BUCKET, + dest_dir=remote_dir, upload_if=gs.UploadIf.IF_MODIFIED) + new_timestamp = gs.get_last_modified_time( + bucket=TEST_BUCKET, path=sample_file_remote_path) + assert old_timestamp != new_timestamp, '%s != %s' % ( + old_timestamp, new_timestamp) + finally: + # Delete all the files we uploaded to Google Storage. + for filename in filenames: + gs.delete_file(bucket=TEST_BUCKET, + path=posixpath.join(remote_dir, subdir, filename)) + finally: + # Clean up the local dir. + shutil.rmtree(local_dir) + + def _test_authenticated_round_trip(): gs = _get_authenticated_gs_handle() remote_dir = _get_unique_posix_dir() @@ -275,7 +355,8 @@ if __name__ == '__main__': - _test_upload_if() + _test_upload_if_multiple_files() + _test_upload_if_one_file() _test_public_read() _test_authenticated_round_trip() _test_dir_upload_and_download()