fix(storage): fix upload object with bucket cmek enabled (#158)

* fix(storage): fix upload object with bucket cmek enabled

* fix(storage): add comments for condition used

* fix(storage): nit

Co-authored-by: Frank Natividad <frankyn@users.noreply.github.com>

google/cloud/storage/blob.py

@@ -1280,7 +1280,15 @@ def _do_multipart_upload(
         if self.user_project is not None:
             name_value_pairs.append(("userProject", self.user_project))
 
-        if self.kms_key_name is not None:
+        # When a Customer Managed Encryption Key is used to encrypt Cloud Storage object
+        # at rest, object resource metadata will store the version of the Key Management
+        # Service cryptographic material. If a Blob instance with KMS Key metadata set is
+        # used to upload a new version of the object then the existing kmsKeyName version
+        # value can't be used in the upload request and the client instead ignores it.
+        if (
+            self.kms_key_name is not None
+            and "cryptoKeyVersions" not in self.kms_key_name
+        ):
             name_value_pairs.append(("kmsKeyName", self.kms_key_name))
 
         if predefined_acl is not None:
@@ -1417,7 +1425,15 @@ def _initiate_resumable_upload(
         if self.user_project is not None:
             name_value_pairs.append(("userProject", self.user_project))
 
-        if self.kms_key_name is not None:
+        # When a Customer Managed Encryption Key is used to encrypt Cloud Storage object
+        # at rest, object resource metadata will store the version of the Key Management
+        # Service cryptographic material. If a Blob instance with KMS Key metadata set is
+        # used to upload a new version of the object then the existing kmsKeyName version
+        # value can't be used in the upload request and the client instead ignores it.
+        if (
+            self.kms_key_name is not None
+            and "cryptoKeyVersions" not in self.kms_key_name
+        ):
             name_value_pairs.append(("kmsKeyName", self.kms_key_name))
 
         if predefined_acl is not None:

tests/system/test_system.py

@@ -1966,6 +1966,30 @@ def test_rewrite_rotate_csek_to_cmek(self):
 
         self.assertEqual(dest.download_as_string(), source_data)
 
+    def test_upload_new_blob_w_bucket_cmek_enabled(self):
+        blob_name = "test-blob"
+        payload = b"DEADBEEF"
+        alt_payload = b"NEWDEADBEEF"
+
+        kms_key_name = self._kms_key_name()
+        self.bucket.default_kms_key_name = kms_key_name
+        self.bucket.patch()
+        self.assertEqual(self.bucket.default_kms_key_name, kms_key_name)
+
+        blob = self.bucket.blob(blob_name)
+        blob.upload_from_string(payload)
+        # We don't know the current version of the key.
+        self.assertTrue(blob.kms_key_name.startswith(kms_key_name))
+
+        blob.upload_from_string(alt_payload, if_generation_match=blob.generation)
+        self.case_blobs_to_delete.append(blob)
+
+        self.assertEqual(blob.download_as_string(), alt_payload)
+
+        self.bucket.default_kms_key_name = None
+        self.bucket.patch()
+        self.assertIsNone(self.bucket.default_kms_key_name)
+
 
 class TestRetentionPolicy(unittest.TestCase):
     def setUp(self):

tests/unit/test_blob.py

@@ -1525,7 +1525,7 @@ def _do_multipart_success(
         if predefined_acl is not None:
             qs_params.append(("predefinedAcl", predefined_acl))
 
-        if kms_key_name is not None:
+        if kms_key_name is not None and "cryptoKeyVersions" not in kms_key_name:
             qs_params.append(("kmsKeyName", kms_key_name))
 
         if if_generation_match is not None:
@@ -1579,6 +1579,17 @@ def test__do_multipart_upload_with_kms(self, mock_get_boundary):
         )
         self._do_multipart_success(mock_get_boundary, kms_key_name=kms_resource)
 
+    @mock.patch(u"google.resumable_media._upload.get_boundary", return_value=b"==0==")
+    def test__do_multipart_upload_with_kms_with_version(self, mock_get_boundary):
+        kms_resource = (
+            "projects/test-project-123/"
+            "locations/us/"
+            "keyRings/test-ring/"
+            "cryptoKeys/test-key"
+            "cryptoKeyVersions/1"
+        )
+        self._do_multipart_success(mock_get_boundary, kms_key_name=kms_resource)
+
     @mock.patch(u"google.resumable_media._upload.get_boundary", return_value=b"==0==")
     def test__do_multipart_upload_with_retry(self, mock_get_boundary):
         self._do_multipart_success(mock_get_boundary, num_retries=8)
@@ -1685,7 +1696,7 @@ def _initiate_resumable_helper(
         if predefined_acl is not None:
             qs_params.append(("predefinedAcl", predefined_acl))
 
-        if kms_key_name is not None:
+        if kms_key_name is not None and "cryptoKeyVersions" not in kms_key_name:
             qs_params.append(("kmsKeyName", kms_key_name))
 
         if if_generation_match is not None:
@@ -1770,6 +1781,16 @@ def test__initiate_resumable_upload_with_kms(self):
         )
         self._initiate_resumable_helper(kms_key_name=kms_resource)
 
+    def test__initiate_resumable_upload_with_kms_with_version(self):
+        kms_resource = (
+            "projects/test-project-123/"
+            "locations/us/"
+            "keyRings/test-ring/"
+            "cryptoKeys/test-key"
+            "cryptoKeyVersions/1"
+        )
+        self._initiate_resumable_helper(kms_key_name=kms_resource)
+
     def test__initiate_resumable_upload_without_chunk_size(self):
         self._initiate_resumable_helper(blob_chunk_size=None)