diff --git a/storage/cloud-client/snippets.py b/storage/cloud-client/snippets.py index c6c69d6ae66..3d490891a4b 100644 --- a/storage/cloud-client/snippets.py +++ b/storage/cloud-client/snippets.py @@ -23,6 +23,7 @@ import argparse import datetime +import pprint from google.cloud import storage @@ -42,6 +43,45 @@ def delete_bucket(bucket_name): print('Bucket {} deleted'.format(bucket.name)) +def get_bucket_labels(bucket_name): + """Prints out a bucket's labels.""" + storage_client = storage.Client() + bucket = storage_client.get_bucket(bucket_name) + labels = bucket.labels + pprint.pprint(labels) + + +def add_bucket_label(bucket_name): + """Add a label to a bucket.""" + storage_client = storage.Client() + bucket = storage_client.get_bucket(bucket_name) + + labels = bucket.labels + labels['example'] = 'label' + bucket.labels = labels + bucket.patch() + + print('Updated labels on {}.'.format(bucket.name)) + pprint.pprint(bucket.labels) + + +def remove_bucket_label(bucket_name): + """Remove a label from a bucket.""" + storage_client = storage.Client() + bucket = storage_client.get_bucket(bucket_name) + + labels = bucket.labels + + if 'example' in labels: + del labels['example'] + + bucket.labels = labels + bucket.patch() + + print('Updated labels on {}.'.format(bucket.name)) + pprint.pprint(bucket.labels) + + def list_blobs(bucket_name): """Lists all the blobs in the bucket.""" storage_client = storage.Client() @@ -222,6 +262,10 @@ def copy_blob(bucket_name, blob_name, new_bucket_name, new_blob_name): subparsers = parser.add_subparsers(dest='command') subparsers.add_parser('create-bucket', help=create_bucket.__doc__) subparsers.add_parser('delete-bucket', help=delete_bucket.__doc__) + subparsers.add_parser('get-bucket-labels', help=get_bucket_labels.__doc__) + subparsers.add_parser('add-bucket-label', help=add_bucket_label.__doc__) + subparsers.add_parser( + 'remove-bucket-label', help=remove_bucket_label.__doc__) subparsers.add_parser('list', help=list_blobs.__doc__) list_with_prefix_parser = subparsers.add_parser( @@ -268,6 +312,12 @@ def copy_blob(bucket_name, blob_name, new_bucket_name, new_blob_name): create_bucket(args.bucket_name) elif args.command == 'delete-bucket': delete_bucket(args.bucket_name) + if args.command == 'get-bucket-labels': + get_bucket_labels(args.bucket_name) + if args.command == 'add-bucket-label': + add_bucket_label(args.bucket_name) + if args.command == 'remove-bucket-label': + remove_bucket_label(args.bucket_name) elif args.command == 'list': list_blobs(args.bucket_name) elif args.command == 'list-with-prefix': diff --git a/storage/cloud-client/snippets_test.py b/storage/cloud-client/snippets_test.py index a44e8ebf1e9..bf5ecb8e1e8 100644 --- a/storage/cloud-client/snippets_test.py +++ b/storage/cloud-client/snippets_test.py @@ -25,6 +25,27 @@ BUCKET = os.environ['CLOUD_STORAGE_BUCKET'] +def test_get_bucket_labels(): + snippets.get_bucket_labels(BUCKET) + + +def test_add_bucket_label(capsys): + snippets.add_bucket_label(BUCKET) + out, _ = capsys.readouterr() + assert 'example' in out + + +@pytest.mark.xfail( + reason=( + 'https://github.com/GoogleCloudPlatform' + '/google-cloud-python/issues/3711')) +def test_remove_bucket_label(capsys): + snippets.add_bucket_label(BUCKET) + snippets.remove_bucket_label(BUCKET) + out, _ = capsys.readouterr() + assert '{}' in out + + @pytest.fixture def test_blob(): """Provides a pre-existing blob in the test bucket."""