Skip to content

Commit

Permalink
feat: add additional_blob_attributes to upload_many_from_filenames (#…
Browse files Browse the repository at this point in the history
…1162)

Fixes #996 🦕
  • Loading branch information
andrewsg committed Oct 12, 2023
1 parent c5a983d commit c7229f2
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 0 deletions.
17 changes: 17 additions & 0 deletions google/cloud/storage/transfer_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -427,6 +427,8 @@ def upload_many_from_filenames(
raise_exception=False,
worker_type=PROCESS,
max_workers=DEFAULT_MAX_WORKERS,
*,
additional_blob_attributes=None,
):
"""Upload many files concurrently by their filenames.
Expand Down Expand Up @@ -557,6 +559,17 @@ def upload_many_from_filenames(
and the default is a conservative number that should work okay in most
cases without consuming excessive resources.
:type additional_blob_attributes: dict
:param additional_blob_attributes:
A dictionary of blob attribute names and values. This allows the
configuration of blobs beyond what is possible with
blob_constructor_kwargs. For instance, {"cache_control": "no-cache"}
would set the cache_control attribute of each blob to "no-cache".
As with blob_constructor_kwargs, this affects the creation of every
blob identically. To fine-tune each blob individually, use `upload_many`
and create the blobs as desired before passing them in.
:raises: :exc:`concurrent.futures.TimeoutError` if deadline is exceeded.
:rtype: list
Expand All @@ -567,13 +580,17 @@ def upload_many_from_filenames(
"""
if blob_constructor_kwargs is None:
blob_constructor_kwargs = {}
if additional_blob_attributes is None:
additional_blob_attributes = {}

file_blob_pairs = []

for filename in filenames:
path = os.path.join(source_directory, filename)
blob_name = blob_name_prefix + filename
blob = bucket.blob(blob_name, **blob_constructor_kwargs)
for prop, value in additional_blob_attributes.items():
setattr(blob, prop, value)
file_blob_pairs.append((path, blob))

return upload_many(
Expand Down
19 changes: 19 additions & 0 deletions tests/system/test_transfer_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,25 @@ def test_upload_many_skip_if_exists(
assert len(blobs_to_delete) == 1


def test_upload_many_from_filenames_with_attributes(
listable_bucket, listable_filenames, file_data, blobs_to_delete
):
SOURCE_DIRECTORY, FILENAME = os.path.split(file_data["logo"]["path"])

transfer_manager.upload_many_from_filenames(
listable_bucket,
[FILENAME],
source_directory=SOURCE_DIRECTORY,
additional_blob_attributes={"cache_control": "no-cache"},
raise_exception=True,
)

blob = listable_bucket.blob(FILENAME)
blob.reload()
blobs_to_delete.append(blob)
assert blob.cache_control == "no-cache"


def test_download_many(listable_bucket):
blobs = list(listable_bucket.list_blobs())
with tempfile.TemporaryDirectory() as tempdir:
Expand Down
32 changes: 32 additions & 0 deletions tests/unit/test_transfer_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -482,6 +482,38 @@ def test_upload_many_from_filenames_minimal_args():
bucket.blob.assert_any_call(FILENAMES[1])


def test_upload_many_from_filenames_additional_properties():
bucket = mock.Mock()
blob = mock.Mock()
bucket_blob = mock.Mock(return_value=blob)
blob.cache_control = None
bucket.blob = bucket_blob

FILENAME = "file_a.txt"
ADDITIONAL_BLOB_ATTRIBUTES = {"cache_control": "no-cache"}
EXPECTED_FILE_BLOB_PAIRS = [(FILENAME, mock.ANY)]

with mock.patch(
"google.cloud.storage.transfer_manager.upload_many"
) as mock_upload_many:
transfer_manager.upload_many_from_filenames(
bucket, [FILENAME], additional_blob_attributes=ADDITIONAL_BLOB_ATTRIBUTES
)

mock_upload_many.assert_called_once_with(
EXPECTED_FILE_BLOB_PAIRS,
skip_if_exists=False,
upload_kwargs=None,
deadline=None,
raise_exception=False,
worker_type=transfer_manager.PROCESS,
max_workers=8,
)

for attrib, value in ADDITIONAL_BLOB_ATTRIBUTES.items():
assert getattr(blob, attrib) == value


def test_download_many_to_path():
bucket = mock.Mock()

Expand Down

0 comments on commit c7229f2

Please sign in to comment.