From 5deba50b8c2d91d08bd5f5fb68742268c494b4a9 Mon Sep 17 00:00:00 2001
From: sriram <sriramsridhar@users.noreply.github.com>
Date: Fri, 15 Sep 2023 22:21:56 +0530
Subject: [PATCH] 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 <swast@google.com>
---
 google/cloud/bigquery/dataset.py | 33 ++++++++++++++++++++++++++++++++
 tests/unit/test_dataset.py       | 25 +++++++++++++++++++++++-
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py
index 513c32d9c..114f0de18 100644
--- a/google/cloud/bigquery/dataset.py
+++ b/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.
diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py
index 5e26a0c03..f2bdf8db5 100644
--- a/tests/unit/test_dataset.py
+++ b/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")