feat: Custom Placement Config Dual Region Support (#819)

cojenco · web-flow · commit febece768022 · 2022-07-13T09:48:14.000-07:00
* refactor: dual-region API update

* test coverage
diff --git a/google/cloud/storage/bucket.py b/google/cloud/storage/bucket.py
@@ -2418,13 +2418,27 @@ def location(self, value):
         warnings.warn(_LOCATION_SETTER_MESSAGE, DeprecationWarning, stacklevel=2)
         self._location = value
 
+    @property
+    def data_locations(self):
+        """Retrieve the list of regional locations for custom dual-region buckets.
+
+        See https://cloud.google.com/storage/docs/json_api/v1/buckets and
+        https://cloud.google.com/storage/docs/locations
+
+        Returns ``None`` if the property has not been set before creation,
+        if the bucket's resource has not been loaded from the server,
+        or if the bucket is not a dual-regions bucket.
+        :rtype: list of str or ``NoneType``
+        """
+        custom_placement_config = self._properties.get("customPlacementConfig", {})
+        return custom_placement_config.get("dataLocations")
+
     @property
     def location_type(self):
-        """Retrieve or set the location type for the bucket.
+        """Retrieve the location type for the bucket.
 
         See https://cloud.google.com/storage/docs/storage-classes
 
-        :setter: Set the location type for this bucket.
         :getter: Gets the the location type for this bucket.
 
         :rtype: str or ``NoneType``
diff --git a/google/cloud/storage/client.py b/google/cloud/storage/client.py
@@ -602,6 +602,7 @@ def _post_resource(
             google.cloud.exceptions.NotFound
                 If the bucket is not found.
         """
+
         return self._connection.api_request(
             method="POST",
             path=path,
@@ -847,6 +848,7 @@ def create_bucket(
         project=None,
         user_project=None,
         location=None,
+        data_locations=None,
         predefined_acl=None,
         predefined_default_object_acl=None,
         timeout=_DEFAULT_TIMEOUT,
@@ -876,7 +878,11 @@ def create_bucket(
             location (str):
                 (Optional) The location of the bucket. If not passed,
                 the default location, US, will be used. If specifying a dual-region,
-                can be specified as a string, e.g., 'US-CENTRAL1+US-WEST1'. See:
+                `data_locations` should be set in conjunction.. See:
+                https://cloud.google.com/storage/docs/locations
+            data_locations (list of str):
+                (Optional) The list of regional locations of a custom dual-region bucket.
+                Dual-regions require exactly 2 regional locations. See:
                 https://cloud.google.com/storage/docs/locations
             predefined_acl (str):
                 (Optional) Name of predefined ACL to apply to bucket. See:
@@ -979,6 +985,9 @@ def create_bucket(
         if location is not None:
             properties["location"] = location
 
+        if data_locations is not None:
+            properties["customPlacementConfig"] = {"dataLocations": data_locations}
+
         api_response = self._post_resource(
             "/b",
             properties,
diff --git a/samples/README.md b/samples/README.md
@@ -324,7 +324,7 @@ View the [source code](https://github.com/googleapis/python-storage/blob/main/sa
 View the [source code](https://github.com/googleapis/python-storage/blob/main/samples/snippets/storage_create_bucket_dual_region.py). To run this sample:
 
 
-`python storage_create_bucket_dual_region.py <BUCKET_NAME> <REGION_1> <REGION_2>`
+`python storage_create_bucket_dual_region.py <BUCKET_NAME> <LOCATION> <REGION_1> <REGION_2>`
 
 -----
 ### Create Bucket Notifications
diff --git a/samples/snippets/snippets_test.py b/samples/snippets/snippets_test.py
@@ -435,10 +435,11 @@ def test_create_bucket_class_location(test_bucket_create):
 
 
 def test_create_bucket_dual_region(test_bucket_create, capsys):
+    location = "US"
     region_1 = "US-EAST1"
     region_2 = "US-WEST1"
     storage_create_bucket_dual_region.create_bucket_dual_region(
-        test_bucket_create.name, region_1, region_2
+        test_bucket_create.name, location, region_1, region_2
     )
     out, _ = capsys.readouterr()
     assert f"Bucket {test_bucket_create.name} created in {region_1}+{region_2}" in out
diff --git a/samples/snippets/storage_create_bucket_dual_region.py b/samples/snippets/storage_create_bucket_dual_region.py
@@ -24,8 +24,8 @@
 from google.cloud import storage
 
 
-def create_bucket_dual_region(bucket_name, region_1, region_2):
-    """Creates a Dual-Region Bucket with provided locations."""
+def create_bucket_dual_region(bucket_name, location, region_1, region_2):
+    """Creates a Dual-Region Bucket with provided location and regions.."""
     # The ID of your GCS bucket
     # bucket_name = "your-bucket-name"
 
@@ -34,9 +34,10 @@ def create_bucket_dual_region(bucket_name, region_1, region_2):
     # https://cloud.google.com/storage/docs/locations
     # region_1 = "US-EAST1"
     # region_2 = "US-WEST1"
+    # location = "US"
 
     storage_client = storage.Client()
-    storage_client.create_bucket(bucket_name, location=f"{region_1}+{region_2}")
+    storage_client.create_bucket(bucket_name, location=location, data_locations=[region_1, region_2])
 
     print(f"Bucket {bucket_name} created in {region_1}+{region_2}.")
 
@@ -46,5 +47,5 @@ def create_bucket_dual_region(bucket_name, region_1, region_2):
 
 if __name__ == "__main__":
     create_bucket_dual_region(
-        bucket_name=sys.argv[1], region_1=sys.argv[2], region_2=sys.argv[3]
+        bucket_name=sys.argv[1], location=sys.argv[2], region_1=sys.argv[3], region_2=sys.argv[4]
     )
diff --git a/tests/system/test_client.py b/tests/system/test_client.py
@@ -68,21 +68,21 @@ def test_create_bucket_dual_region(storage_client, buckets_to_delete):
     from google.cloud.storage.constants import DUAL_REGION_LOCATION_TYPE
 
     new_bucket_name = _helpers.unique_name("dual-region-bucket")
-    region_1 = "US-EAST1"
-    region_2 = "US-WEST1"
-    dual_region = f"{region_1}+{region_2}"
+    location = "US"
+    data_locations = ["US-EAST1", "US-WEST1"]
 
     with pytest.raises(exceptions.NotFound):
         storage_client.get_bucket(new_bucket_name)
 
     created = _helpers.retry_429_503(storage_client.create_bucket)(
-        new_bucket_name, location=dual_region
+        new_bucket_name, location=location, data_locations=data_locations
     )
     buckets_to_delete.append(created)
 
     assert created.name == new_bucket_name
-    assert created.location == dual_region
+    assert created.location == location
     assert created.location_type == DUAL_REGION_LOCATION_TYPE
+    assert created.data_locations == data_locations
 
 
 def test_list_buckets(storage_client, buckets_to_delete):
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -1451,6 +1451,44 @@ def test_create_bucket_w_explicit_location(self):
             _target_object=bucket,
         )
 
+    def test_create_bucket_w_custom_dual_region(self):
+        project = "PROJECT"
+        bucket_name = "bucket-name"
+        location = "US"
+        data_locations = ["US-EAST1", "US-WEST1"]
+        api_response = {
+            "location": location,
+            "customPlacementConfig": {"dataLocations": data_locations},
+            "name": bucket_name,
+        }
+        credentials = _make_credentials()
+        client = self._make_one(project=project, credentials=credentials)
+        client._post_resource = mock.Mock()
+        client._post_resource.return_value = api_response
+
+        bucket = client.create_bucket(
+            bucket_name, location=location, data_locations=data_locations
+        )
+
+        self.assertEqual(bucket.location, location)
+        self.assertEqual(bucket.data_locations, data_locations)
+
+        expected_path = "/b"
+        expected_data = {
+            "location": location,
+            "customPlacementConfig": {"dataLocations": data_locations},
+            "name": bucket_name,
+        }
+        expected_query_params = {"project": project}
+        client._post_resource.assert_called_once_with(
+            expected_path,
+            expected_data,
+            query_params=expected_query_params,
+            timeout=self._get_default_timeout(),
+            retry=DEFAULT_RETRY,
+            _target_object=bucket,
+        )
+
     def test_create_bucket_w_explicit_project(self):
         project = "PROJECT"
         other_project = "other-project-123"
