Repository: googleapis/google-cloud-python

Issue:
Storage: blob.update_storage_class() does not work for large files
[`Blob.update_storage_class()`](https://googleapis.github.io/google-cloud-python/latest/_modules/google/cloud/storage/blob.html#Blob.update_storage_class) does not work the same as [`Blob.rewrite()`](https://googleapis.github.io/google-cloud-python/latest/_modules/google/cloud/storage/blob.html#Blob.rewrite) despite sharing a common underlying API. When multiple API calls are required to complete the rewrite operation, `Blob.update_storage_class()` will fail after the first call. `Blob.update_storage_class()` assumes that `api_response['resource']` exists. However, according to the [docs](https://cloud.google.com/storage/docs/json_api/v1/objects/rewrite#response), this is only true when the operation completes on the first operation. Instead, I think that `Blob.update_storage_class()` should just leverage the logic in `Blob.rewrite()`, which handles this case already. I will also look to see how `gsutil` implements this.

#### Stack trace
```
Traceback (most recent call last):
  File "rehost.py", line 31, in <module>
    blob.update_storage_class('COLDLINE')
  File "/var/lib/mysql/.local/lib/python2.7/site-packages/google/cloud/storage/blob.py", line 1497, in update_storage_class
    self._set_properties(api_response['resource'])
KeyError: 'resource'
Command exited with non-zero status 1
```

#### Implementation

- [x] Re-implement `Blob.update_storage_class()` as a wrapper around `Blob.rewrite()`
- [x] Update any `Blob.update_storage_class()` tests
- [x] Manually test new implementation with a large file

#### References
How `gsutil` handles looping through a rewrite operation - https://github.com/GoogleCloudPlatform/gsutil/blob/704781b250314791dda08191a7472d1a72da2c6c/gslib/gcs_json_api.py#L1477-L1526


Discussion:
I plan to contribute this myself.

Resolving patch:
diff --git a/storage/google/cloud/storage/blob.py b/storage/google/cloud/storage/blob.py
--- a/storage/google/cloud/storage/blob.py
+++ b/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
 


Test patch:
diff --git a/storage/tests/system.py b/storage/tests/system.py
--- a/storage/tests/system.py
+++ b/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
diff --git a/storage/tests/unit/test_blob.py b/storage/tests/unit/test_blob.py
--- a/storage/tests/unit/test_blob.py
+++ b/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)