Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: implementing Storage Billing model property to BQ Datasets #1643

Merged
merged 10 commits into from Sep 15, 2023
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