Skip to content

Commit

Permalink
feat: add Dataset.storage_billing_model setter, use `client.update_…
Browse files Browse the repository at this point in the history
…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>
  • Loading branch information
sriramsridhar and tswast committed Sep 15, 2023
1 parent 7248f1f commit 5deba50
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 1 deletion.
33 changes: 33 additions & 0 deletions google/cloud/bigquery/dataset.py
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down
25 changes: 24 additions & 1 deletion tests/unit/test_dataset.py
Expand Up @@ -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},
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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")
Expand Down

0 comments on commit 5deba50

Please sign in to comment.