Update Blob.update_storage_class to support rewrite tokens (#6527)

storage/google/cloud/storage/blob.py

@@ -1512,7 +1512,9 @@ def rewrite(self, source, token=None, client=None):
         return api_response["rewriteToken"], rewritten, size
 
     def update_storage_class(self, new_class, client=None):
-        """Update blob's storage class via a rewrite-in-place.
+        """Update blob's storage class via a rewrite-in-place. This helper will
+        wait for the rewrite to complete before returning, so it may take some
+        time for large files.
 
         See
         https://cloud.google.com/storage/docs/per-object-storage-class
@@ -1530,25 +1532,13 @@ def update_storage_class(self, new_class, client=None):
         if new_class not in self._STORAGE_CLASSES:
             raise ValueError("Invalid storage class: %s" % (new_class,))
 
-        client = self._require_client(client)
-
-        query_params = {}
-
-        if self.user_project is not None:
-            query_params["userProject"] = self.user_project
-
-        headers = _get_encryption_headers(self._encryption_key)
-        headers.update(_get_encryption_headers(self._encryption_key, source=True))
+        # Update current blob's storage class prior to rewrite
+        self._patch_property('storageClass', new_class)
 
-        api_response = client._connection.api_request(
-            method="POST",
-            path=self.path + "/rewriteTo" + self.path,
-            query_params=query_params,
-            data={"storageClass": new_class},
-            headers=headers,
-            _target_object=self,
-        )
-        self._set_properties(api_response["resource"])
+        # Execute consecutive rewrite operations until operation is done
+        token, _, _ = self.rewrite(self)
+        while token is not None:
+            token, _, _ = self.rewrite(self, token=token)
 
     cache_control = _scalar_property("cacheControl")
     """HTTP 'Cache-Control' header for this object.
@@ -1815,7 +1805,7 @@ def kms_key_name(self):
     This can only be set at blob / object **creation** time. If you'd
     like to change the storage class **after** the blob / object already
     exists in a bucket, call :meth:`update_storage_class` (which uses
-    the "storage.objects.rewrite" method).
+    :meth:`rewrite`).
 
     See https://cloud.google.com/storage/docs/storage-classes
 

storage/tests/system.py

@@ -963,6 +963,39 @@ def test_rewrite_rotate_with_user_project(self):
             retry_429(created.delete)(force=True)
 
 
+class TestStorageUpdateStorageClass(TestStorageFiles):
+
+    def test_update_storage_class_small_file(self):
+        blob = self.bucket.blob("SmallFile")
+
+        file_data = self.FILES["simple"]
+        blob.upload_from_filename(file_data["path"])
+        self.case_blobs_to_delete.append(blob)
+
+        blob.update_storage_class("NEARLINE")
+        blob.reload()
+        self.assertEqual(blob.storage_class, "NEARLINE")
+
+        blob.update_storage_class("COLDLINE")
+        blob.reload()
+        self.assertEqual(blob.storage_class, "COLDLINE")
+
+    def test_update_storage_class_large_file(self):
+        blob = self.bucket.blob("BigFile")
+
+        file_data = self.FILES["big"]
+        blob.upload_from_filename(file_data["path"])
+        self.case_blobs_to_delete.append(blob)
+
+        blob.update_storage_class("NEARLINE")
+        blob.reload()
+        self.assertEqual(blob.storage_class, "NEARLINE")
+
+        blob.update_storage_class("COLDLINE")
+        blob.reload()
+        self.assertEqual(blob.storage_class, "COLDLINE")
+
+
 class TestStorageNotificationCRUD(unittest.TestCase):
 
     topic = None

storage/tests/unit/test_blob.py

@@ -2499,10 +2499,43 @@ def test_update_storage_class_invalid(self):
         with self.assertRaises(ValueError):
             blob.update_storage_class(u"BOGUS")
 
+    def test_update_storage_class_large_file(self):
+        BLOB_NAME = "blob-name"
+        STORAGE_CLASS = u"NEARLINE"
+        TOKEN = "TOKEN"
+        INCOMPLETE_RESPONSE = {
+            "totalBytesRewritten": 42,
+            "objectSize": 84,
+            "done": False,
+            "rewriteToken": TOKEN,
+            "resource": {"storageClass": STORAGE_CLASS}
+        }
+        COMPLETE_RESPONSE = {
+            "totalBytesRewritten": 84,
+            "objectSize": 84,
+            "done": True,
+            "resource": {"storageClass": STORAGE_CLASS}
+        }
+        response_1 = ({"status": http_client.OK}, INCOMPLETE_RESPONSE)
+        response_2 = ({"status": http_client.OK}, COMPLETE_RESPONSE)
+        connection = _Connection(response_1, response_2)
+        client = _Client(connection)
+        bucket = _Bucket(client=client)
+        blob = self._make_one(BLOB_NAME, bucket=bucket)
+
+        blob.update_storage_class("NEARLINE")
+
+        self.assertEqual(blob.storage_class, "NEARLINE")
+
     def test_update_storage_class_wo_encryption_key(self):
         BLOB_NAME = "blob-name"
         STORAGE_CLASS = u"NEARLINE"
-        RESPONSE = {"resource": {"storageClass": STORAGE_CLASS}}
+        RESPONSE = {
+            "totalBytesRewritten": 42,
+            "objectSize": 42,
+            "done": True,
+            "resource": {"storageClass": STORAGE_CLASS}
+        }
         response = ({"status": http_client.OK}, RESPONSE)
         connection = _Connection(response)
         client = _Client(connection)
@@ -2542,7 +2575,12 @@ def test_update_storage_class_w_encryption_key_w_user_project(self):
         BLOB_KEY_HASH_B64 = base64.b64encode(BLOB_KEY_HASH).rstrip().decode("ascii")
         STORAGE_CLASS = u"NEARLINE"
         USER_PROJECT = "user-project-123"
-        RESPONSE = {"resource": {"storageClass": STORAGE_CLASS}}
+        RESPONSE = {
+            "totalBytesRewritten": 42,
+            "objectSize": 42,
+            "done": True,
+            "resource": {"storageClass": STORAGE_CLASS}
+        }
         response = ({"status": http_client.OK}, RESPONSE)
         connection = _Connection(response)
         client = _Client(connection)