Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat: add preconditions and retry configuration to blob.create_resuma…
…ble_upload_session (#484)

* feat: add preconditions and retry configuration to blob.create_resumable_upload_session

* move imports
  • Loading branch information
cojenco committed Jun 30, 2021
1 parent c027ccf commit 0ae35ee
Show file tree
Hide file tree
Showing 2 changed files with 109 additions and 7 deletions.
58 changes: 58 additions & 0 deletions google/cloud/storage/blob.py
Expand Up @@ -2782,6 +2782,11 @@ def create_resumable_upload_session(
client=None,
timeout=_DEFAULT_TIMEOUT,
checksum=None,
if_generation_match=None,
if_generation_not_match=None,
if_metageneration_match=None,
if_metageneration_not_match=None,
retry=DEFAULT_RETRY_IF_GENERATION_SPECIFIED,
):
"""Create a resumable upload session.
Expand Down Expand Up @@ -2857,6 +2862,41 @@ def create_resumable_upload_session(
delete the uploaded object automatically. Supported values
are "md5", "crc32c" and None. The default is None.
:type if_generation_match: long
:param if_generation_match:
(Optional) See :ref:`using-if-generation-match`
:type if_generation_not_match: long
:param if_generation_not_match:
(Optional) See :ref:`using-if-generation-not-match`
:type if_metageneration_match: long
:param if_metageneration_match:
(Optional) See :ref:`using-if-metageneration-match`
:type if_metageneration_not_match: long
:param if_metageneration_not_match:
(Optional) See :ref:`using-if-metageneration-not-match`
:type retry: google.api_core.retry.Retry or google.cloud.storage.retry.ConditionalRetryPolicy
:param retry: (Optional) How to retry the RPC. A None value will disable
retries. A google.api_core.retry.Retry value will enable retries,
and the object will define retriable response codes and errors and
configure backoff and timeout options.
A google.cloud.storage.retry.ConditionalRetryPolicy value wraps a
Retry object and activates it only if certain conditions are met.
This class exists to provide safe defaults for RPC calls that are
not technically safe to retry normally (due to potential data
duplication or other side-effects) but become safe to retry if a
condition such as if_generation_match is set.
See the retry.py source code and docstrings in this package
(google.cloud.storage.retry) for information on retry types and how
to configure them.
Media operations (downloads and uploads) do not support non-default
predicates in a Retry object. The default will always be used. Other
configuration changes for Retry objects such as delays and deadlines
are respected.
:rtype: str
:returns: The resumable upload session URL. The upload can be
completed by making an HTTP PUT request with the
Expand All @@ -2865,6 +2905,19 @@ def create_resumable_upload_session(
:raises: :class:`google.cloud.exceptions.GoogleCloudError`
if the session creation response returns an error status.
"""

# Handle ConditionalRetryPolicy.
if isinstance(retry, ConditionalRetryPolicy):
# Conditional retries are designed for non-media calls, which change
# arguments into query_params dictionaries. Media operations work
# differently, so here we make a "fake" query_params to feed to the
# ConditionalRetryPolicy.
query_params = {
"ifGenerationMatch": if_generation_match,
"ifMetagenerationMatch": if_metageneration_match,
}
retry = retry.get_retry_policy_if_conditions_met(query_params=query_params)

extra_headers = {}
if origin is not None:
# This header is specifically for client-side uploads, it
Expand All @@ -2883,10 +2936,15 @@ def create_resumable_upload_session(
size,
None,
predefined_acl=None,
if_generation_match=if_generation_match,
if_generation_not_match=if_generation_not_match,
if_metageneration_match=if_metageneration_match,
if_metageneration_not_match=if_metageneration_not_match,
extra_headers=extra_headers,
chunk_size=self._CHUNK_SIZE_MULTIPLE,
timeout=timeout,
checksum=checksum,
retry=retry,
)

return upload.resumable_url
Expand Down
58 changes: 51 additions & 7 deletions tests/unit/test_blob.py
Expand Up @@ -25,6 +25,7 @@
import pytest
import six
from six.moves import http_client
from six.moves.urllib.parse import urlencode

from google.cloud.storage.retry import DEFAULT_RETRY
from google.cloud.storage.retry import DEFAULT_RETRY_IF_ETAG_IN_JSON
Expand Down Expand Up @@ -2041,8 +2042,6 @@ def _do_multipart_success(
mtls=False,
retry=None,
):
from six.moves.urllib.parse import urlencode

bucket = _Bucket(name="w00t", user_project=user_project)
blob = self._make_one(u"blob-name", bucket=bucket, kms_key_name=kms_key_name)
self.assertIsNone(blob.chunk_size)
Expand Down Expand Up @@ -2286,7 +2285,6 @@ def _initiate_resumable_helper(
mtls=False,
retry=None,
):
from six.moves.urllib.parse import urlencode
from google.resumable_media.requests import ResumableUpload
from google.cloud.storage.blob import _DEFAULT_CHUNKSIZE

Expand Down Expand Up @@ -3248,7 +3246,15 @@ def test_upload_from_string_w_text_w_num_retries(self):
self._upload_from_string_helper(data, num_retries=2)

def _create_resumable_upload_session_helper(
self, origin=None, side_effect=None, timeout=None
self,
origin=None,
side_effect=None,
timeout=None,
if_generation_match=None,
if_generation_not_match=None,
if_metageneration_match=None,
if_metageneration_not_match=None,
retry=None,
):
bucket = _Bucket(name="alex-trebek")
blob = self._make_one("blob-name", bucket=bucket)
Expand Down Expand Up @@ -3280,6 +3286,11 @@ def _create_resumable_upload_session_helper(
size=size,
origin=origin,
client=client,
if_generation_match=if_generation_match,
if_generation_not_match=if_generation_not_match,
if_metageneration_match=if_metageneration_match,
if_metageneration_not_match=if_metageneration_not_match,
retry=retry,
**timeout_kwarg
)

Expand All @@ -3289,10 +3300,23 @@ def _create_resumable_upload_session_helper(

# Check the mocks.
upload_url = (
"https://storage.googleapis.com/upload/storage/v1"
+ bucket.path
+ "/o?uploadType=resumable"
"https://storage.googleapis.com/upload/storage/v1" + bucket.path + "/o"
)

qs_params = [("uploadType", "resumable")]
if if_generation_match is not None:
qs_params.append(("ifGenerationMatch", if_generation_match))

if if_generation_not_match is not None:
qs_params.append(("ifGenerationNotMatch", if_generation_not_match))

if if_metageneration_match is not None:
qs_params.append(("ifMetagenerationMatch", if_metageneration_match))

if if_metageneration_not_match is not None:
qs_params.append(("ifMetaGenerationNotMatch", if_metageneration_not_match))

upload_url += "?" + urlencode(qs_params)
payload = b'{"name": "blob-name"}'
expected_headers = {
"content-type": "application/json; charset=UTF-8",
Expand All @@ -3318,6 +3342,26 @@ def test_create_resumable_upload_session_with_custom_timeout(self):
def test_create_resumable_upload_session_with_origin(self):
self._create_resumable_upload_session_helper(origin="http://google.com")

def test_create_resumable_upload_session_with_generation_match(self):
self._create_resumable_upload_session_helper(
if_generation_match=123456, if_metageneration_match=2
)

def test_create_resumable_upload_session_with_generation_not_match(self):
self._create_resumable_upload_session_helper(
if_generation_not_match=0, if_metageneration_not_match=3
)

def test_create_resumable_upload_session_with_conditional_retry_success(self):
self._create_resumable_upload_session_helper(
retry=DEFAULT_RETRY_IF_GENERATION_SPECIFIED, if_generation_match=123456
)

def test_create_resumable_upload_session_with_conditional_retry_failure(self):
self._create_resumable_upload_session_helper(
retry=DEFAULT_RETRY_IF_GENERATION_SPECIFIED
)

def test_create_resumable_upload_session_with_failure(self):
from google.resumable_media import InvalidResponse
from google.cloud import exceptions
Expand Down

0 comments on commit 0ae35ee

Please sign in to comment.