feat: add preconditions and retry configuration to blob.create_resumable_upload_session (#484)

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

* move imports

Author: cojenco

google/cloud/storage/blob.py

@@ -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.
 
@@ -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
@@ -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
@@ -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

tests/unit/test_blob.py

@@ -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
@@ -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)
@@ -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
 
@@ -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)
@@ -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
         )
 
@@ -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",
@@ -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