feat: support includeFoldersAsPrefixes (#1223)

cojenco · web-flow · commit 7bb806538cf3 · 2024-03-18T16:18:33.000-07:00
* feat: support includeFoldersAsPrefixes

* sys test

* update sys test with cleanup
diff --git a/google/cloud/storage/bucket.py b/google/cloud/storage/bucket.py
@@ -1306,6 +1306,7 @@ def list_blobs(
         timeout=_DEFAULT_TIMEOUT,
         retry=DEFAULT_RETRY,
         match_glob=None,
+        include_folders_as_prefixes=None,
         soft_deleted=None,
     ):
         """Return an iterator used to find blobs in the bucket.
@@ -1388,6 +1389,11 @@ def list_blobs(
             The string value must be UTF-8 encoded. See:
             https://cloud.google.com/storage/docs/json_api/v1/objects/list#list-object-glob
 
+        :type include_folders_as_prefixes: bool
+            (Optional) If true, includes Folders and Managed Folders in the set of
+            ``prefixes`` returned by the query. Only applicable if ``delimiter`` is set to /.
+            See: https://cloud.google.com/storage/docs/managed-folders
+
         :type soft_deleted: bool
         :param soft_deleted:
             (Optional) If true, only soft-deleted objects will be listed as distinct results in order of increasing
@@ -1415,6 +1421,7 @@ def list_blobs(
             timeout=timeout,
             retry=retry,
             match_glob=match_glob,
+            include_folders_as_prefixes=include_folders_as_prefixes,
             soft_deleted=soft_deleted,
         )
 
diff --git a/google/cloud/storage/client.py b/google/cloud/storage/client.py
@@ -1184,6 +1184,7 @@ def list_blobs(
         timeout=_DEFAULT_TIMEOUT,
         retry=DEFAULT_RETRY,
         match_glob=None,
+        include_folders_as_prefixes=None,
         soft_deleted=None,
     ):
         """Return an iterator used to find blobs in the bucket.
@@ -1283,6 +1284,11 @@ def list_blobs(
                 The string value must be UTF-8 encoded. See:
                 https://cloud.google.com/storage/docs/json_api/v1/objects/list#list-object-glob
 
+            include_folders_as_prefixes (bool):
+                (Optional) If true, includes Folders and Managed Folders in the set of
+                ``prefixes`` returned by the query. Only applicable if ``delimiter`` is set to /.
+                See: https://cloud.google.com/storage/docs/managed-folders
+
             soft_deleted (bool):
                 (Optional) If true, only soft-deleted objects will be listed as distinct results in order of increasing
                 generation number. This parameter can only be used successfully if the bucket has a soft delete policy.
@@ -1325,6 +1331,9 @@ def list_blobs(
         if fields is not None:
             extra_params["fields"] = fields
 
+        if include_folders_as_prefixes is not None:
+            extra_params["includeFoldersAsPrefixes"] = include_folders_as_prefixes
+
         if soft_deleted is not None:
             extra_params["softDeleted"] = soft_deleted
 
diff --git a/tests/system/test_bucket.py b/tests/system/test_bucket.py
@@ -653,6 +653,47 @@ def test_bucket_list_blobs_w_match_glob(
         assert [blob.name for blob in blobs] == expected_names
 
 
+def test_bucket_list_blobs_include_managed_folders(
+    storage_client,
+    buckets_to_delete,
+    blobs_to_delete,
+    hierarchy_filenames,
+):
+    bucket_name = _helpers.unique_name("ubla-mf")
+    bucket = storage_client.bucket(bucket_name)
+    bucket.iam_configuration.uniform_bucket_level_access_enabled = True
+    _helpers.retry_429_503(bucket.create)()
+    buckets_to_delete.append(bucket)
+
+    payload = b"helloworld"
+    for filename in hierarchy_filenames:
+        blob = bucket.blob(filename)
+        blob.upload_from_string(payload)
+        blobs_to_delete.append(blob)
+
+    # Make API call to create a managed folder.
+    # TODO: change to use storage control client once available.
+    path = f"/b/{bucket_name}/managedFolders"
+    properties = {"name": "managedfolder1"}
+    storage_client._post_resource(path, properties)
+
+    expected_prefixes = set(["parent/"])
+    blob_iter = bucket.list_blobs(delimiter="/")
+    list(blob_iter)
+    assert blob_iter.prefixes == expected_prefixes
+
+    # Test that managed folders are only included when IncludeFoldersAsPrefixes is set.
+    expected_prefixes = set(["parent/", "managedfolder1/"])
+    blob_iter = bucket.list_blobs(delimiter="/", include_folders_as_prefixes=True)
+    list(blob_iter)
+    assert blob_iter.prefixes == expected_prefixes
+
+    # Cleanup: API call to delete a managed folder.
+    # TODO: change to use storage control client once available.
+    path = f"/b/{bucket_name}/managedFolders/managedfolder1"
+    storage_client._delete_resource(path)
+
+
 def test_bucket_update_retention_period(
     storage_client,
     buckets_to_delete,
diff --git a/tests/unit/test_bucket.py b/tests/unit/test_bucket.py
@@ -1177,6 +1177,7 @@ def test_list_blobs_w_defaults(self):
         expected_versions = None
         expected_projection = "noAcl"
         expected_fields = None
+        expected_include_folders_as_prefixes = None
         soft_deleted = None
         client.list_blobs.assert_called_once_with(
             bucket,
@@ -1193,6 +1194,7 @@ def test_list_blobs_w_defaults(self):
             timeout=self._get_default_timeout(),
             retry=DEFAULT_RETRY,
             match_glob=expected_match_glob,
+            include_folders_as_prefixes=expected_include_folders_as_prefixes,
             soft_deleted=soft_deleted,
         )
 
@@ -1206,6 +1208,7 @@ def test_list_blobs_w_explicit(self):
         start_offset = "c"
         end_offset = "g"
         include_trailing_delimiter = True
+        include_folders_as_prefixes = True
         versions = True
         soft_deleted = True
         projection = "full"
@@ -1231,6 +1234,7 @@ def test_list_blobs_w_explicit(self):
             timeout=timeout,
             retry=retry,
             match_glob=match_glob,
+            include_folders_as_prefixes=include_folders_as_prefixes,
             soft_deleted=soft_deleted,
         )
 
@@ -1247,6 +1251,7 @@ def test_list_blobs_w_explicit(self):
         expected_versions = versions
         expected_projection = projection
         expected_fields = fields
+        expected_include_folders_as_prefixes = include_folders_as_prefixes
         expected_soft_deleted = soft_deleted
         other_client.list_blobs.assert_called_once_with(
             bucket,
@@ -1263,6 +1268,7 @@ def test_list_blobs_w_explicit(self):
             timeout=timeout,
             retry=retry,
             match_glob=expected_match_glob,
+            include_folders_as_prefixes=expected_include_folders_as_prefixes,
             soft_deleted=expected_soft_deleted,
         )
 
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -2015,6 +2015,7 @@ def test_list_blobs_w_explicit_w_user_project(self):
         start_offset = "c"
         end_offset = "g"
         include_trailing_delimiter = True
+        include_folders_as_prefixes = True
         soft_deleted = False
         versions = True
         projection = "full"
@@ -2048,6 +2049,7 @@ def test_list_blobs_w_explicit_w_user_project(self):
             timeout=timeout,
             retry=retry,
             match_glob=match_glob,
+            include_folders_as_prefixes=include_folders_as_prefixes,
             soft_deleted=soft_deleted,
         )
 
@@ -2070,6 +2072,7 @@ def test_list_blobs_w_explicit_w_user_project(self):
             "versions": versions,
             "fields": fields,
             "userProject": user_project,
+            "includeFoldersAsPrefixes": include_folders_as_prefixes,
             "softDeleted": soft_deleted,
         }
         expected_page_start = _blobs_page_start
