feat: add *generation*match args into Blob.compose() (#122)

Gurov Ilya · frankyn · web-flow · commit dc01c59e0361 · 2020-06-09T13:29:31.000-07:00
* feat: add *generation*match args into Blob.compose() * add test case with error * new compose surface * Revert "new compose surface" This reverts commit 2ddda40. * add an error for different length args, add a usage example * add condition to avoid sending params with None * specify comments Co-authored-by: Frank Natividad <frankyn@users.noreply.github.com>
diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py
@@ -2234,34 +2234,104 @@ def make_private(self, client=None):
         self.acl.all().revoke_read()
         self.acl.save(client=client)
 
-    def compose(self, sources, client=None, timeout=_DEFAULT_TIMEOUT):
+    def compose(
+        self,
+        sources,
+        client=None,
+        timeout=_DEFAULT_TIMEOUT,
+        if_generation_match=None,
+        if_metageneration_match=None,
+    ):
         """Concatenate source blobs into this one.
 
         If :attr:`user_project` is set on the bucket, bills the API request
         to that project.
 
         :type sources: list of :class:`Blob`
-        :param sources: blobs whose contents will be composed into this blob.
+        :param sources: Blobs whose contents will be composed into this blob.
 
         :type client: :class:`~google.cloud.storage.client.Client` or
                       ``NoneType``
-        :param client: (Optional) The client to use.  If not passed, falls back
+        :param client: (Optional) The client to use. If not passed, falls back
                        to the ``client`` stored on the blob's bucket.
+
         :type timeout: float or tuple
         :param timeout: (Optional) The amount of time, in seconds, to wait
             for the server response.
 
             Can also be passed as a tuple (connect_timeout, read_timeout).
             See :meth:`requests.Session.request` documentation for details.
+
+        :type if_generation_match: list of long
+        :param if_generation_match: (Optional) Make the operation conditional on whether
+                                    the blob's current generation matches the given value.
+                                    Setting to 0 makes the operation succeed only if there
+                                    are no live versions of the blob. The list must match
+                                    ``sources`` item-to-item.
+
+        :type if_metageneration_match: list of long
+        :param if_metageneration_match: (Optional) Make the operation conditional on whether
+                                        the blob's current metageneration matches the given
+                                        value. The list must match ``sources`` item-to-item.
+
+        Example:
+            Compose blobs using generation match preconditions.
+
+            >>> from google.cloud import storage
+            >>> client = storage.Client()
+            >>> bucket = client.bucket("bucket-name")
+
+            >>> blobs = [bucket.blob("blob-name-1"), bucket.blob("blob-name-2")]
+            >>> if_generation_match = [None] * len(blobs)
+            >>> if_generation_match[0] = "123"  # precondition for "blob-name-1"
+
+            >>> composed_blob = bucket.blob("composed-name")
+            >>> composed_blob.compose(blobs, if_generation_match)
         """
+        sources_len = len(sources)
+        if if_generation_match is not None and len(if_generation_match) != sources_len:
+            raise ValueError(
+                "'if_generation_match' length must be the same as 'sources' length"
+            )
+
+        if (
+            if_metageneration_match is not None
+            and len(if_metageneration_match) != sources_len
+        ):
+            raise ValueError(
+                "'if_metageneration_match' length must be the same as 'sources' length"
+            )
+
         client = self._require_client(client)
         query_params = {}
 
         if self.user_project is not None:
             query_params["userProject"] = self.user_project
 
+        source_objects = []
+        for index, source in enumerate(sources):
+            source_object = {"name": source.name}
+
+            preconditions = {}
+            if (
+                if_generation_match is not None
+                and if_generation_match[index] is not None
+            ):
+                preconditions["ifGenerationMatch"] = if_generation_match[index]
+
+            if (
+                if_metageneration_match is not None
+                and if_metageneration_match[index] is not None
+            ):
+                preconditions["ifMetagenerationMatch"] = if_metageneration_match[index]
+
+            if preconditions:
+                source_object["objectPreconditions"] = preconditions
+
+            source_objects.append(source_object)
+
         request = {
-            "sourceObjects": [{"name": source.name} for source in sources],
+            "sourceObjects": source_objects,
             "destination": self._properties.copy(),
         }
         api_response = client._connection.api_request(
diff --git a/tests/system/test_system.py b/tests/system/test_system.py
@@ -1441,6 +1441,34 @@ def test_compose_replace_existing_blob(self):
         composed = original.download_as_string()
         self.assertEqual(composed, BEFORE + TO_APPEND)
 
+    def test_compose_with_generation_match(self):
+        BEFORE = b"AAA\n"
+        original = self.bucket.blob("original")
+        original.content_type = "text/plain"
+        original.upload_from_string(BEFORE)
+        self.case_blobs_to_delete.append(original)
+
+        TO_APPEND = b"BBB\n"
+        to_append = self.bucket.blob("to_append")
+        to_append.upload_from_string(TO_APPEND)
+        self.case_blobs_to_delete.append(to_append)
+
+        with self.assertRaises(google.api_core.exceptions.PreconditionFailed):
+            original.compose(
+                [original, to_append],
+                if_generation_match=[6, 7],
+                if_metageneration_match=[8, 9],
+            )
+
+        original.compose(
+            [original, to_append],
+            if_generation_match=[original.generation, to_append.generation],
+            if_metageneration_match=[original.metageneration, to_append.metageneration],
+        )
+
+        composed = original.download_as_string()
+        self.assertEqual(composed, BEFORE + TO_APPEND)
+
     @unittest.skipUnless(USER_PROJECT, "USER_PROJECT not set in environment.")
     def test_compose_with_user_project(self):
         new_bucket_name = "compose-user-project" + unique_resource_id("-")
diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py
@@ -2676,7 +2676,7 @@ def test_make_private(self):
     def test_compose_wo_content_type_set(self):
         SOURCE_1 = "source-1"
         SOURCE_2 = "source-2"
-        DESTINATION = "destinaton"
+        DESTINATION = "destination"
         RESOURCE = {}
         after = ({"status": http_client.OK}, RESOURCE)
         connection = _Connection(after)
@@ -2711,7 +2711,7 @@ def test_compose_wo_content_type_set(self):
     def test_compose_minimal_w_user_project(self):
         SOURCE_1 = "source-1"
         SOURCE_2 = "source-2"
-        DESTINATION = "destinaton"
+        DESTINATION = "destination"
         RESOURCE = {"etag": "DEADBEEF"}
         USER_PROJECT = "user-project-123"
         after = ({"status": http_client.OK}, RESOURCE)
@@ -2747,7 +2747,7 @@ def test_compose_minimal_w_user_project(self):
     def test_compose_w_additional_property_changes(self):
         SOURCE_1 = "source-1"
         SOURCE_2 = "source-2"
-        DESTINATION = "destinaton"
+        DESTINATION = "destination"
         RESOURCE = {"etag": "DEADBEEF"}
         after = ({"status": http_client.OK}, RESOURCE)
         connection = _Connection(after)
@@ -2785,6 +2785,129 @@ def test_compose_w_additional_property_changes(self):
             },
         )
 
+    def test_compose_w_generation_match(self):
+        SOURCE_1 = "source-1"
+        SOURCE_2 = "source-2"
+        DESTINATION = "destination"
+        RESOURCE = {}
+        GENERATION_NUMBERS = [6, 9]
+        METAGENERATION_NUMBERS = [7, 1]
+
+        after = ({"status": http_client.OK}, RESOURCE)
+        connection = _Connection(after)
+        client = _Client(connection)
+        bucket = _Bucket(client=client)
+        source_1 = self._make_one(SOURCE_1, bucket=bucket)
+        source_2 = self._make_one(SOURCE_2, bucket=bucket)
+
+        destination = self._make_one(DESTINATION, bucket=bucket)
+        destination.compose(
+            sources=[source_1, source_2],
+            if_generation_match=GENERATION_NUMBERS,
+            if_metageneration_match=METAGENERATION_NUMBERS,
+        )
+
+        kw = connection._requested
+        self.assertEqual(len(kw), 1)
+        self.assertEqual(
+            kw[0],
+            {
+                "method": "POST",
+                "path": "/b/name/o/%s/compose" % DESTINATION,
+                "query_params": {},
+                "data": {
+                    "sourceObjects": [
+                        {
+                            "name": source_1.name,
+                            "objectPreconditions": {
+                                "ifGenerationMatch": GENERATION_NUMBERS[0],
+                                "ifMetagenerationMatch": METAGENERATION_NUMBERS[0],
+                            },
+                        },
+                        {
+                            "name": source_2.name,
+                            "objectPreconditions": {
+                                "ifGenerationMatch": GENERATION_NUMBERS[1],
+                                "ifMetagenerationMatch": METAGENERATION_NUMBERS[1],
+                            },
+                        },
+                    ],
+                    "destination": {},
+                },
+                "_target_object": destination,
+                "timeout": self._get_default_timeout(),
+            },
+        )
+
+    def test_compose_w_generation_match_bad_length(self):
+        SOURCE_1 = "source-1"
+        SOURCE_2 = "source-2"
+        DESTINATION = "destination"
+        GENERATION_NUMBERS = [6]
+        METAGENERATION_NUMBERS = [7]
+
+        after = ({"status": http_client.OK}, {})
+        connection = _Connection(after)
+        client = _Client(connection)
+        bucket = _Bucket(client=client)
+        source_1 = self._make_one(SOURCE_1, bucket=bucket)
+        source_2 = self._make_one(SOURCE_2, bucket=bucket)
+
+        destination = self._make_one(DESTINATION, bucket=bucket)
+
+        with self.assertRaises(ValueError):
+            destination.compose(
+                sources=[source_1, source_2], if_generation_match=GENERATION_NUMBERS,
+            )
+        with self.assertRaises(ValueError):
+            destination.compose(
+                sources=[source_1, source_2],
+                if_metageneration_match=METAGENERATION_NUMBERS,
+            )
+
+    def test_compose_w_generation_match_nones(self):
+        SOURCE_1 = "source-1"
+        SOURCE_2 = "source-2"
+        DESTINATION = "destination"
+        GENERATION_NUMBERS = [6, None]
+
+        after = ({"status": http_client.OK}, {})
+        connection = _Connection(after)
+        client = _Client(connection)
+        bucket = _Bucket(client=client)
+        source_1 = self._make_one(SOURCE_1, bucket=bucket)
+        source_2 = self._make_one(SOURCE_2, bucket=bucket)
+
+        destination = self._make_one(DESTINATION, bucket=bucket)
+        destination.compose(
+            sources=[source_1, source_2], if_generation_match=GENERATION_NUMBERS,
+        )
+
+        kw = connection._requested
+        self.assertEqual(len(kw), 1)
+        self.assertEqual(
+            kw[0],
+            {
+                "method": "POST",
+                "path": "/b/name/o/%s/compose" % DESTINATION,
+                "query_params": {},
+                "data": {
+                    "sourceObjects": [
+                        {
+                            "name": source_1.name,
+                            "objectPreconditions": {
+                                "ifGenerationMatch": GENERATION_NUMBERS[0],
+                            },
+                        },
+                        {"name": source_2.name},
+                    ],
+                    "destination": {},
+                },
+                "_target_object": destination,
+                "timeout": self._get_default_timeout(),
+            },
+        )
+
     def test_rewrite_response_without_resource(self):
         SOURCE_BLOB = "source"
         DEST_BLOB = "dest"
