feat: add support for restore token (#1369)

cojenco · web-flow · commit 06ed15b33dc8 · 2024-10-30T14:17:43.000-07:00
* feat: add support for restore token

* add unit tests coverage

* update docstrings

* fix docs
diff --git a/google/cloud/storage/_helpers.py b/google/cloud/storage/_helpers.py
@@ -226,6 +226,7 @@ def reload(
         timeout=_DEFAULT_TIMEOUT,
         retry=DEFAULT_RETRY,
         soft_deleted=None,
+        restore_token=None,
     ):
         """Reload properties from Cloud Storage.
 
@@ -278,6 +279,13 @@ def reload(
             the object metadata if the object exists and is in a soft-deleted state.
             :attr:`generation` is required to be set on the blob if ``soft_deleted`` is set to True.
             See: https://cloud.google.com/storage/docs/soft-delete
+
+        :type restore_token: str
+        :param restore_token:
+            (Optional) The restore_token is required to retrieve a soft-deleted object only if
+            its name and generation value do not uniquely identify it, and hierarchical namespace
+            is enabled on the bucket. Otherwise, this parameter is optional.
+            See: https://cloud.google.com/storage/docs/json_api/v1/objects/get
         """
         client = self._require_client(client)
         query_params = self._query_params
@@ -296,6 +304,8 @@ def reload(
             # Soft delete reload requires a generation, even for targets
             # that don't include them in default query params (buckets).
             query_params["generation"] = self.generation
+        if restore_token is not None:
+            query_params["restoreToken"] = restore_token
         headers = self._encryption_headers()
         _add_etag_match_headers(
             headers, if_etag_match=if_etag_match, if_etag_not_match=if_etag_not_match
diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py
@@ -653,6 +653,7 @@ def exists(
         timeout=_DEFAULT_TIMEOUT,
         retry=DEFAULT_RETRY,
         soft_deleted=None,
+        restore_token=None,
     ):
         """Determines whether or not this blob exists.
 
@@ -704,6 +705,13 @@ def exists(
             :attr:`generation` is required to be set on the blob if ``soft_deleted`` is set to True.
             See: https://cloud.google.com/storage/docs/soft-delete
 
+        :type restore_token: str
+        :param restore_token:
+            (Optional) The restore_token is required to retrieve a soft-deleted object only if
+            its name and generation value do not uniquely identify it, and hierarchical namespace
+            is enabled on the bucket. Otherwise, this parameter is optional.
+            See: https://cloud.google.com/storage/docs/json_api/v1/objects/get
+
         :rtype: bool
         :returns: True if the blob exists in Cloud Storage.
         """
@@ -714,6 +722,8 @@ def exists(
         query_params["fields"] = "name"
         if soft_deleted is not None:
             query_params["softDeleted"] = soft_deleted
+        if restore_token is not None:
+            query_params["restoreToken"] = restore_token
 
         _add_generation_match_parameters(
             query_params,
@@ -4794,6 +4804,19 @@ def hard_delete_time(self):
         if hard_delete_time is not None:
             return _rfc3339_nanos_to_datetime(hard_delete_time)
 
+    @property
+    def restore_token(self):
+        """The restore token, a universally unique identifier (UUID), along with the object's
+        name and generation value, uniquely identifies a soft-deleted object.
+        This field is only returned for soft-deleted objects in hierarchical namespace buckets.
+
+        :rtype: string or ``NoneType``
+        :returns:
+            (readonly) The restore token used to differentiate soft-deleted objects with the same name and generation.
+            This field is only returned for soft-deleted objects in hierarchical namespace buckets.
+        """
+        return self._properties.get("restoreToken")
+
 
 def _get_host_name(connection):
     """Returns the host name from the given connection.
diff --git a/google/cloud/storage/bucket.py b/google/cloud/storage/bucket.py
@@ -1256,6 +1256,7 @@ def get_blob(
         timeout=_DEFAULT_TIMEOUT,
         retry=DEFAULT_RETRY,
         soft_deleted=None,
+        restore_token=None,
         **kwargs,
     ):
         """Get a blob object by name.
@@ -1323,6 +1324,13 @@ def get_blob(
             Object ``generation`` is required if ``soft_deleted`` is set to True.
             See: https://cloud.google.com/storage/docs/soft-delete
 
+        :type restore_token: str
+        :param restore_token:
+            (Optional) The restore_token is required to retrieve a soft-deleted object only if
+            its name and generation value do not uniquely identify it, and hierarchical namespace
+            is enabled on the bucket. Otherwise, this parameter is optional.
+            See: https://cloud.google.com/storage/docs/json_api/v1/objects/get
+
         :param kwargs: Keyword arguments to pass to the
                        :class:`~google.cloud.storage.blob.Blob` constructor.
 
@@ -1351,6 +1359,7 @@ def get_blob(
                 if_metageneration_not_match=if_metageneration_not_match,
                 retry=retry,
                 soft_deleted=soft_deleted,
+                restore_token=restore_token,
             )
         except NotFound:
             return None
@@ -2199,6 +2208,7 @@ def restore_blob(
         generation=None,
         copy_source_acl=None,
         projection=None,
+        restore_token=None,
         if_generation_match=None,
         if_generation_not_match=None,
         if_metageneration_match=None,
@@ -2229,6 +2239,13 @@ def restore_blob(
         :param projection: (Optional) Specifies the set of properties to return.
                            If used, must be 'full' or 'noAcl'.
 
+        :type restore_token: str
+        :param restore_token:
+            (Optional) The restore_token is required to restore a soft-deleted object
+            only if its name and generation value do not uniquely identify it, and hierarchical namespace
+            is enabled on the bucket. Otherwise, this parameter is optional.
+            See: https://cloud.google.com/storage/docs/json_api/v1/objects/restore
+
         :type if_generation_match: long
         :param if_generation_match:
             (Optional) See :ref:`using-if-generation-match`
@@ -2276,6 +2293,8 @@ def restore_blob(
             query_params["copySourceAcl"] = copy_source_acl
         if projection is not None:
             query_params["projection"] = projection
+        if restore_token is not None:
+            query_params["restoreToken"] = restore_token
 
         _add_generation_match_parameters(
             query_params,
diff --git a/tests/system/test_bucket.py b/tests/system/test_bucket.py
@@ -1232,7 +1232,7 @@ def test_soft_delete_policy(
     assert restored_blob.generation != gen
 
     # Patch the soft delete policy on an existing bucket.
-    new_duration_secs = 10 * 86400
+    new_duration_secs = 0
     bucket.soft_delete_policy.retention_duration_seconds = new_duration_secs
     bucket.patch()
     assert bucket.soft_delete_policy.retention_duration_seconds == new_duration_secs
@@ -1265,3 +1265,55 @@ def test_new_bucket_with_hierarchical_namespace(
     bucket = storage_client.create_bucket(bucket_obj)
     buckets_to_delete.append(bucket)
     assert bucket.hierarchical_namespace_enabled is True
+
+
+def test_restore_token(
+    storage_client,
+    buckets_to_delete,
+    blobs_to_delete,
+):
+    # Create HNS bucket with soft delete policy.
+    duration_secs = 7 * 86400
+    bucket = storage_client.bucket(_helpers.unique_name("w-soft-delete"))
+    bucket.hierarchical_namespace_enabled = True
+    bucket.iam_configuration.uniform_bucket_level_access_enabled = True
+    bucket.soft_delete_policy.retention_duration_seconds = duration_secs
+    bucket = _helpers.retry_429_503(storage_client.create_bucket)(bucket)
+    buckets_to_delete.append(bucket)
+
+    # Insert an object and delete it to enter soft-deleted state.
+    payload = b"DEADBEEF"
+    blob_name = _helpers.unique_name("soft-delete")
+    blob = bucket.blob(blob_name)
+    blob.upload_from_string(payload)
+    # blob = bucket.get_blob(blob_name)
+    gen = blob.generation
+    blob.delete()
+
+    # Get the soft-deleted object and restore token.
+    blob = bucket.get_blob(blob_name, generation=gen, soft_deleted=True)
+    restore_token = blob.restore_token
+
+    # List and get soft-deleted object that includes restore token.
+    all_blobs = list(bucket.list_blobs(soft_deleted=True))
+    assert all_blobs[0].restore_token is not None
+    blob_w_restore_token = bucket.get_blob(
+        blob_name, generation=gen, soft_deleted=True, restore_token=restore_token
+    )
+    assert blob_w_restore_token.soft_delete_time is not None
+    assert blob_w_restore_token.hard_delete_time is not None
+    assert blob_w_restore_token.restore_token is not None
+
+    # Restore the soft-deleted object using the restore token.
+    restored_blob = bucket.restore_blob(
+        blob_name, generation=gen, restore_token=restore_token
+    )
+    blobs_to_delete.append(restored_blob)
+    assert restored_blob.exists() is True
+    assert restored_blob.generation != gen
+
+    # Patch the soft delete policy on the bucket.
+    new_duration_secs = 0
+    bucket.soft_delete_policy.retention_duration_seconds = new_duration_secs
+    bucket.patch()
+    assert bucket.soft_delete_policy.retention_duration_seconds == new_duration_secs
diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py
@@ -784,21 +784,25 @@ def test_exists_hit_w_generation_w_retry(self):
             _target_object=None,
         )
 
-    def test_exists_hit_w_generation_w_soft_deleted(self):
+    def test_exists_hit_w_gen_soft_deleted_restore_token(self):
         blob_name = "blob-name"
         generation = 123456
+        restore_token = "88ba0d97-639e-5902"
         api_response = {"name": blob_name}
         client = mock.Mock(spec=["_get_resource"])
         client._get_resource.return_value = api_response
         bucket = _Bucket(client)
         blob = self._make_one(blob_name, bucket=bucket, generation=generation)
 
-        self.assertTrue(blob.exists(retry=None, soft_deleted=True))
+        self.assertTrue(
+            blob.exists(retry=None, soft_deleted=True, restore_token=restore_token)
+        )
 
         expected_query_params = {
             "fields": "name",
             "generation": generation,
             "softDeleted": True,
+            "restoreToken": restore_token,
         }
         expected_headers = {}
         client._get_resource.assert_called_once_with(
@@ -5870,6 +5874,16 @@ def test_soft_hard_delete_time_getter(self):
         self.assertEqual(blob.soft_delete_time, soft_timstamp)
         self.assertEqual(blob.hard_delete_time, hard_timstamp)
 
+    def test_restore_token_getter(self):
+        BLOB_NAME = "blob-name"
+        bucket = _Bucket()
+        restore_token = "88ba0d97-639e-5902"
+        properties = {
+            "restoreToken": restore_token,
+        }
+        blob = self._make_one(BLOB_NAME, bucket=bucket, properties=properties)
+        self.assertEqual(blob.restore_token, restore_token)
+
     def test_soft_hard_delte_time_unset(self):
         BUCKET = object()
         blob = self._make_one("blob-name", bucket=BUCKET)
diff --git a/tests/unit/test_bucket.py b/tests/unit/test_bucket.py
@@ -1018,18 +1018,24 @@ def test_get_blob_hit_w_user_project(self):
             _target_object=blob,
         )
 
-    def test_get_blob_hit_w_generation_w_soft_deleted(self):
+    def test_get_blob_hit_w_gen_soft_deleted_restore_token(self):
         from google.cloud.storage.blob import Blob
 
         name = "name"
         blob_name = "blob-name"
         generation = 1512565576797178
+        restore_token = "88ba0d97-639e-5902"
         api_response = {"name": blob_name, "generation": generation}
         client = mock.Mock(spec=["_get_resource"])
         client._get_resource.return_value = api_response
         bucket = self._make_one(client, name=name)
 
-        blob = bucket.get_blob(blob_name, generation=generation, soft_deleted=True)
+        blob = bucket.get_blob(
+            blob_name,
+            generation=generation,
+            soft_deleted=True,
+            restore_token=restore_token,
+        )
 
         self.assertIsInstance(blob, Blob)
         self.assertIs(blob.bucket, bucket)
@@ -1041,6 +1047,7 @@ def test_get_blob_hit_w_generation_w_soft_deleted(self):
             "generation": generation,
             "projection": "noAcl",
             "softDeleted": True,
+            "restoreToken": restore_token,
         }
         expected_headers = {}
         client._get_resource.assert_called_once_with(
@@ -4217,8 +4224,10 @@ def test_restore_blob_w_explicit(self):
         user_project = "user-project-123"
         bucket_name = "restore_bucket"
         blob_name = "restore_blob"
+        new_generation = 987655
         generation = 123456
-        api_response = {"name": blob_name, "generation": generation}
+        restore_token = "88ba0d97-639e-5902"
+        api_response = {"name": blob_name, "generation": new_generation}
         client = mock.Mock(spec=["_post_resource"])
         client._post_resource.return_value = api_response
         bucket = self._make_one(
@@ -4233,6 +4242,8 @@ def test_restore_blob_w_explicit(self):
         restored_blob = bucket.restore_blob(
             blob_name,
             client=client,
+            generation=generation,
+            restore_token=restore_token,
             if_generation_match=if_generation_match,
             if_generation_not_match=if_generation_not_match,
             if_metageneration_match=if_metageneration_match,
@@ -4245,6 +4256,8 @@ def test_restore_blob_w_explicit(self):
         expected_path = f"/b/{bucket_name}/o/{blob_name}/restore"
         expected_data = None
         expected_query_params = {
+            "generation": generation,
+            "restoreToken": restore_token,
             "userProject": user_project,
             "projection": projection,
             "ifGenerationMatch": if_generation_match,
