feat: add Dataset.storage_billing_model setter, use client.update_dataset(ds, fields=["storage_billing_model"]) to update (#1643)

Adding Storage Billing Model property. See: https://cloud.google.com/bigquery/docs/updating-datasets#update_storage_billing_models

---------

Co-authored-by: Tim Swast <[email protected]>

Author: sriramsridhar

google/cloud/bigquery/dataset.py

@@ -527,6 +527,7 @@ class Dataset(object):
         "default_table_expiration_ms": "defaultTableExpirationMs",
         "friendly_name": "friendlyName",
         "default_encryption_configuration": "defaultEncryptionConfiguration",
+        "storage_billing_model": "storageBillingModel",
     }
 
     def __init__(self, dataset_ref) -> None:
@@ -763,6 +764,38 @@ def default_encryption_configuration(self, value):
             api_repr = value.to_api_repr()
         self._properties["defaultEncryptionConfiguration"] = api_repr
 
+    @property
+    def storage_billing_model(self):
+        """Union[str, None]: StorageBillingModel of the dataset as set by the user
+        (defaults to :data:`None`).
+
+        Set the value to one of ``'LOGICAL'`` or ``'PHYSICAL'``. This change
+        takes 24 hours to take effect and you must wait 14 days before you can
+        change the storage billing model again.
+
+        See `storage billing model
+        <https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#Dataset.FIELDS.storage_billing_model>`_
+        in REST API docs and `updating the storage billing model
+        <https://cloud.google.com/bigquery/docs/updating-datasets#update_storage_billing_models>`_
+        guide.
+
+        Raises:
+            ValueError: for invalid value types.
+        """
+        return self._properties.get("storageBillingModel")
+
+    @storage_billing_model.setter
+    def storage_billing_model(self, value):
+        if not isinstance(value, str) and value is not None:
+            raise ValueError(
+                "storage_billing_model must be a string (e.g. 'LOGICAL', 'PHYSICAL'), or None. "
+                f"Got {repr(value)}."
+            )
+        if value:
+            self._properties["storageBillingModel"] = value
+        if value is None:
+            self._properties["storageBillingModel"] = "LOGICAL"
+
     @classmethod
     def from_string(cls, full_dataset_id: str) -> "Dataset":
         """Construct a dataset from fully-qualified dataset ID.

tests/unit/test_dataset.py

@@ -667,6 +667,7 @@ def _make_resource(self):
             "location": "US",
             "selfLink": self.RESOURCE_URL,
             "defaultTableExpirationMs": 3600,
+            "storageBillingModel": "LOGICAL",
             "access": [
                 {"role": "OWNER", "userByEmail": USER_EMAIL},
                 {"role": "OWNER", "groupByEmail": GROUP_EMAIL},
@@ -736,7 +737,12 @@ def _verify_resource_properties(self, dataset, resource):
             )
         else:
             self.assertIsNone(dataset.default_encryption_configuration)
-
+        if "storageBillingModel" in resource:
+            self.assertEqual(
+                dataset.storage_billing_model, resource.get("storageBillingModel")
+            )
+        else:
+            self.assertIsNone(dataset.storage_billing_model)
         if "access" in resource:
             self._verify_access_entry(dataset.access_entries, resource)
         else:
@@ -941,6 +947,23 @@ def test_default_encryption_configuration_setter(self):
         dataset.default_encryption_configuration = None
         self.assertIsNone(dataset.default_encryption_configuration)
 
+    def test_storage_billing_model_setter(self):
+        dataset = self._make_one(self.DS_REF)
+        dataset.storage_billing_model = "PHYSICAL"
+        self.assertEqual(dataset.storage_billing_model, "PHYSICAL")
+
+    def test_storage_billing_model_setter_with_none(self):
+        dataset = self._make_one(self.DS_REF)
+        dataset.storage_billing_model = None
+        self.assertEqual(dataset.storage_billing_model, "LOGICAL")
+
+    def test_storage_billing_model_setter_with_invalid_type(self):
+        dataset = self._make_one(self.DS_REF)
+        with self.assertRaises(ValueError) as raises:
+            dataset.storage_billing_model = object()
+
+        self.assertIn("storage_billing_model", str(raises.exception))
+
     def test_from_string(self):
         cls = self._get_target_class()
         got = cls.from_string("string-project.string_dataset")