make gs_utils.py class-based, so credentials file can be read just once BUG=skia:2618 R=rmistry@google.com Review URL: https://codereview.chromium.org/385013004
diff --git a/py/utils/gs_utils.py b/py/utils/gs_utils.py index c30295b..46c9bf1 100755 --- a/py/utils/gs_utils.py +++ b/py/utils/gs_utils.py
@@ -40,88 +40,110 @@ from boto.s3.prefix import Prefix -def delete_file(bucket, path): - """Delete a single file within a GS bucket. +class GSUtils(object): + """Utilities for accessing Google Cloud Storage, using the boto library.""" - TODO(epoger): what if bucket or path does not exist? Should probably raise - an exception. Implement, and add a test to exercise this. + def __init__(self, boto_file_path=os.path.join('~','.boto')): + """Constructor. - Params: - bucket: GS bucket to delete a file from - path: full path (Posix-style) of the file within the bucket to delete - """ - conn = _create_connection() - b = conn.get_bucket(bucket_name=bucket) - item = Key(b) - item.key = path - item.delete() + Params: + boto_file_path: full path (local-OS-style) on local disk where .boto + credentials file can be found. An exception is thrown if this file + is missing. + TODO(epoger): Change missing-file behavior: allow the caller to + operate on public files in Google Storage. + """ + boto_file_path = os.path.expanduser(boto_file_path) + print 'Reading boto file from %s' % boto_file_path + boto_dict = _config_file_as_dict(filepath=boto_file_path) + self._gs_access_key_id = boto_dict['gs_access_key_id'] + self._gs_secret_access_key = boto_dict['gs_secret_access_key'] + def delete_file(self, bucket, path): + """Delete a single file within a GS bucket. -def upload_file(source_path, dest_bucket, dest_path): - """Upload contents of a local file to Google Storage. + TODO(epoger): what if bucket or path does not exist? Should probably raise + an exception. Implement, and add a test to exercise this. - TODO(epoger): Add the extra parameters provided by upload_file() within - https://github.com/google/skia-buildbot/blob/master/slave/skia_slave_scripts/utils/old_gs_utils.py , - so we can replace that function with this one. + Params: + bucket: GS bucket to delete a file from + path: full path (Posix-style) of the file within the bucket to delete + """ + conn = self._create_connection() + b = conn.get_bucket(bucket_name=bucket) + item = Key(b) + item.key = path + item.delete() - params: - source_path: full path (local-OS-style) on local disk to read from - dest_bucket: GCS bucket to copy the file to - dest_path: full path (Posix-style) within that bucket - """ - conn = _create_connection() - b = conn.get_bucket(bucket_name=dest_bucket) - item = Key(b) - item.key = dest_path - item.set_contents_from_filename(filename=source_path) + def upload_file(self, source_path, dest_bucket, dest_path): + """Upload contents of a local file to Google Storage. + TODO(epoger): Add the extra parameters provided by upload_file() within + https://github.com/google/skia-buildbot/blob/master/slave/skia_slave_scripts/utils/old_gs_utils.py , + so we can replace that function with this one. -def download_file(source_bucket, source_path, dest_path, - create_subdirs_if_needed=False): - """ Downloads a single file from Google Cloud Storage to local disk. + params: + source_path: full path (local-OS-style) on local disk to read from + dest_bucket: GCS bucket to copy the file to + dest_path: full path (Posix-style) within that bucket + """ + conn = self._create_connection() + b = conn.get_bucket(bucket_name=dest_bucket) + item = Key(b) + item.key = dest_path + item.set_contents_from_filename(filename=source_path) - Args: - source_bucket: GCS bucket to download the file from - source_path: full path (Posix-style) within that bucket - dest_path: full path (local-OS-style) on local disk to copy the file to - create_subdirs_if_needed: boolean; whether to create subdirectories as - needed to create dest_path - """ - conn = _create_connection() - b = conn.get_bucket(bucket_name=source_bucket) - item = Key(b) - item.key = source_path - if create_subdirs_if_needed: - _makedirs_if_needed(os.path.dirname(dest_path)) - with open(dest_path, 'w') as f: - item.get_contents_to_file(fp=f) + def download_file(self, source_bucket, source_path, dest_path, + create_subdirs_if_needed=False): + """Downloads a single file from Google Cloud Storage to local disk. + Args: + source_bucket: GCS bucket to download the file from + source_path: full path (Posix-style) within that bucket + dest_path: full path (local-OS-style) on local disk to copy the file to + create_subdirs_if_needed: boolean; whether to create subdirectories as + needed to create dest_path + """ + conn = self._create_connection() + b = conn.get_bucket(bucket_name=source_bucket) + item = Key(b) + item.key = source_path + if create_subdirs_if_needed: + _makedirs_if_needed(os.path.dirname(dest_path)) + with open(dest_path, 'w') as f: + item.get_contents_to_file(fp=f) -def list_bucket_contents(bucket, subdir=None): - """ Returns files in the Google Cloud Storage bucket as a (dirs, files) tuple. + def list_bucket_contents(self, bucket, subdir=None): + """Returns files in the Google Storage bucket as a (dirs, files) tuple. - Args: - bucket: name of the Google Storage bucket - subdir: directory within the bucket to list, or None for root directory - """ - # The GS command relies on the prefix (if any) ending with a slash. - prefix = subdir or '' - if prefix and not prefix.endswith('/'): - prefix += '/' - prefix_length = len(prefix) if prefix else 0 + Args: + bucket: name of the Google Storage bucket + subdir: directory within the bucket to list, or None for root directory + """ + # The GS command relies on the prefix (if any) ending with a slash. + prefix = subdir or '' + if prefix and not prefix.endswith('/'): + prefix += '/' + prefix_length = len(prefix) if prefix else 0 - conn = _create_connection() - b = conn.get_bucket(bucket_name=bucket) - lister = BucketListResultSet(bucket=b, prefix=prefix, delimiter='/') - dirs = [] - files = [] - for item in lister: - t = type(item) - if t is Key: - files.append(item.key[prefix_length:]) - elif t is Prefix: - dirs.append(item.name[prefix_length:-1]) - return (dirs, files) + conn = self._create_connection() + b = conn.get_bucket(bucket_name=bucket) + lister = BucketListResultSet(bucket=b, prefix=prefix, delimiter='/') + dirs = [] + files = [] + for item in lister: + t = type(item) + if t is Key: + files.append(item.key[prefix_length:]) + elif t is Prefix: + dirs.append(item.name[prefix_length:-1]) + return (dirs, files) + + def _create_connection(self): + """Returns a GSConnection object we can use to access Google Storage.""" + return GSConnection( + gs_access_key_id=self._gs_access_key_id, + gs_secret_access_key=self._gs_secret_access_key) def _config_file_as_dict(filepath): @@ -148,29 +170,8 @@ return dic -def _create_connection(boto_file_path=os.path.join('~','.boto')): - """Returns a GSConnection object we can use to access Google Storage. - - Params: - boto_file_path: full path (local-OS-style) on local disk where .boto - credentials file can be found - - TODO(epoger): Change this module to be object-based, where __init__() reads - the boto file into boto_dict once instead of repeatedly for each operation. - - TODO(epoger): if the file does not exist, rather than raising an exception, - create a GSConnection that can operate on public files. - """ - boto_file_path = os.path.expanduser(boto_file_path) - print 'Reading boto file from %s' % boto_file_path - boto_dict = _config_file_as_dict(filepath=boto_file_path) - return GSConnection( - gs_access_key_id=boto_dict['gs_access_key_id'], - gs_secret_access_key=boto_dict['gs_secret_access_key']) - - def _makedirs_if_needed(path): - """ Creates a directory (and any parent directories needed), if it does not + """Creates a directory (and any parent directories needed), if it does not exist yet. Args: @@ -188,6 +189,7 @@ remote_dir = 'gs_utils_test/%d' % random.randint(0, sys.maxint) subdir = 'subdir' filenames_to_upload = ['file1', 'file2'] + gs = GSUtils() # Upload test files to Google Storage. local_src_dir = tempfile.mkdtemp() @@ -196,46 +198,47 @@ for filename in filenames_to_upload: with open(os.path.join(local_src_dir, subdir, filename), 'w') as f: f.write('contents of %s\n' % filename) - upload_file(source_path=os.path.join(local_src_dir, subdir, filename), - dest_bucket=bucket, - dest_path=posixpath.join(remote_dir, subdir, filename)) + gs.upload_file(source_path=os.path.join(local_src_dir, subdir, filename), + dest_bucket=bucket, + dest_path=posixpath.join(remote_dir, subdir, filename)) finally: shutil.rmtree(local_src_dir) # Get a list of the files we uploaded to Google Storage. - (dirs, files) = list_bucket_contents( + (dirs, files) = gs.list_bucket_contents( bucket=bucket, subdir=remote_dir) - assert dirs == [subdir] - assert files == [] - (dirs, files) = list_bucket_contents( + assert dirs == [subdir], '%s == [%s]' % (dirs, subdir) + assert files == [], '%s == []' % files + (dirs, files) = gs.list_bucket_contents( bucket=bucket, subdir=posixpath.join(remote_dir, subdir)) - assert dirs == [] - assert files == filenames_to_upload + assert dirs == [], '%s == []' % dirs + assert files == filenames_to_upload, '%s == %s' % (files, filenames_to_upload) # Download the files we uploaded to Google Storage, and validate contents. local_dest_dir = tempfile.mkdtemp() try: for filename in filenames_to_upload: - download_file(source_bucket=bucket, - source_path=posixpath.join(remote_dir, subdir, filename), - dest_path=os.path.join(local_dest_dir, subdir, filename), - create_subdirs_if_needed=True) + gs.download_file(source_bucket=bucket, + source_path=posixpath.join(remote_dir, subdir, filename), + dest_path=os.path.join(local_dest_dir, subdir, filename), + create_subdirs_if_needed=True) with open(os.path.join(local_dest_dir, subdir, filename)) as f: file_contents = f.read() - assert file_contents == 'contents of %s\n' % filename + assert file_contents == 'contents of %s\n' % filename, ( + '%s == "contents of %s\n"' % (file_contents, filename)) finally: shutil.rmtree(local_dest_dir) # Delete all the files we uploaded to Google Storage. for filename in filenames_to_upload: - delete_file(bucket=bucket, - path=posixpath.join(remote_dir, subdir, filename)) + gs.delete_file(bucket=bucket, + path=posixpath.join(remote_dir, subdir, filename)) # Confirm that we deleted all the files we uploaded to Google Storage. - (dirs, files) = list_bucket_contents( + (dirs, files) = gs.list_bucket_contents( bucket=bucket, subdir=posixpath.join(remote_dir, subdir)) - assert dirs == [] - assert files == [] + assert dirs == [], '%s == []' % dirs + assert files == [], '%s == []' % files # TODO(epoger): How should we exercise this self-test?