From 9cd45974412a3473ecb3cabc241c269c08d14db7 Mon Sep 17 00:00:00 2001 From: Sriram Date: Wed, 16 Aug 2023 13:25:00 +0530 Subject: [PATCH 1/4] Adding Storage Billing model property --- google/cloud/bigquery/dataset.py | 21 +++++++++++++++++++++ tests/unit/test_dataset.py | 17 ++++++++++++++++- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py index 513c32d9c..558eca366 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,26 @@ 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`). + + 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("Pass a string, or None") + if value: + api_repr = value.to_api_repr() + self._properties["storageBillingModel"] = api_repr + 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..05b3a883d 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,15 @@ 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" + ) + dataset.default_encryption_configuration = None + self.assertIsNone(dataset.default_encryption_configuration) + def test_from_string(self): cls = self._get_target_class() got = cls.from_string("string-project.string_dataset") From 7ec5bec6554e559612242639badd4713300d0204 Mon Sep 17 00:00:00 2001 From: Sriram Date: Wed, 16 Aug 2023 13:25:00 +0530 Subject: [PATCH 2/4] feat: implementing Storage Billing model property to BQ Datasets --- google/cloud/bigquery/dataset.py | 21 +++++++++++++++++++++ tests/unit/test_dataset.py | 17 ++++++++++++++++- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py index 513c32d9c..558eca366 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,26 @@ 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`). + + 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("Pass a string, or None") + if value: + api_repr = value.to_api_repr() + self._properties["storageBillingModel"] = api_repr + 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..05b3a883d 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,15 @@ 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" + ) + dataset.default_encryption_configuration = None + self.assertIsNone(dataset.default_encryption_configuration) + def test_from_string(self): cls = self._get_target_class() got = cls.from_string("string-project.string_dataset") From b60c42114f533153b8e189393fc888527f42f9e8 Mon Sep 17 00:00:00 2001 From: Sriram Date: Wed, 16 Aug 2023 18:14:09 +0530 Subject: [PATCH 3/4] removed to_api_repr --- google/cloud/bigquery/dataset.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py index 558eca366..4e23885cc 100644 --- a/google/cloud/bigquery/dataset.py +++ b/google/cloud/bigquery/dataset.py @@ -779,8 +779,7 @@ def storage_billing_model(self, value): if not isinstance(value, str) and value is not None: raise ValueError("Pass a string, or None") if value: - api_repr = value.to_api_repr() - self._properties["storageBillingModel"] = api_repr + self._properties["storageBillingModel"] = value if value is None: self._properties["storageBillingModel"] = "LOGICAL" From 36627de9778ff5a494f35070f943d61776fc41d6 Mon Sep 17 00:00:00 2001 From: Tim Swast Date: Fri, 15 Sep 2023 11:19:29 -0500 Subject: [PATCH 4/4] improve storage_billing_model unit test coverage --- google/cloud/bigquery/dataset.py | 15 ++++++++++++++- tests/unit/test_dataset.py | 18 +++++++++++++----- 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py index 4e23885cc..114f0de18 100644 --- a/google/cloud/bigquery/dataset.py +++ b/google/cloud/bigquery/dataset.py @@ -769,6 +769,16 @@ 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 + `_ + in REST API docs and `updating the storage billing model + `_ + guide. + Raises: ValueError: for invalid value types. """ @@ -777,7 +787,10 @@ def storage_billing_model(self): @storage_billing_model.setter def storage_billing_model(self, value): if not isinstance(value, str) and value is not None: - raise ValueError("Pass a string, or 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: diff --git a/tests/unit/test_dataset.py b/tests/unit/test_dataset.py index 05b3a883d..f2bdf8db5 100644 --- a/tests/unit/test_dataset.py +++ b/tests/unit/test_dataset.py @@ -950,11 +950,19 @@ def test_default_encryption_configuration_setter(self): 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" - ) - dataset.default_encryption_configuration = None - self.assertIsNone(dataset.default_encryption_configuration) + 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()