feat: add support for dataset.default_rounding_mode

Author: Gaurang033

google/cloud/bigquery/dataset.py

@@ -527,7 +527,7 @@ class Dataset(object):
         "friendly_name": "friendlyName",
         "default_encryption_configuration": "defaultEncryptionConfiguration",
         "storage_billing_model": "storageBillingModel",
-        "max_time_travel_hours": "maxTimeTravelHours",
+        "default_rounding_mode": "defaultRoundingMode",
     }
 
     def __init__(self, dataset_ref) -> None:
@@ -536,25 +536,42 @@ def __init__(self, dataset_ref) -> None:
         self._properties = {"datasetReference": dataset_ref.to_api_repr(), "labels": {}}
 
     @property
-    def max_time_travel_hours(self):
-        """
-        Optional[int]: Defines the time travel window in hours. The value can be from 48 to 168 hours (2 to 7 days).
-        value must in multiple of 24 hours (48, 72, 96, 120, 144, 168)
-        The default value is 168 hours if this is not set.
-        """
-        return self._properties.get("maxTimeTravelHours")
+    def default_rounding_mode(self):
+        """Union[str, None]: defaultRoundingMode of the dataset as set by the user
+        (defaults to :data:`None`).
+
+        Set the value to one of ``'ROUND_HALF_AWAY_FROM_ZERO'``, ``'ROUND_HALF_EVEN'``, or
+        ``'ROUNDING_MODE_UNSPECIFIED'``.
 
-    @max_time_travel_hours.setter
-    def max_time_travel_hours(self, hours):
-        if not isinstance(hours, int):
-            raise ValueError("Pass the int value")
-        if hours < 2 * 24 or hours > 7 * 24:
+        See `default rounding mode
+        <https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#Dataset.FIELDS.default_rounding_mode>`_
+        in REST API docs and `updating the default rounding model
+        <https://cloud.google.com/bigquery/docs/updating-datasets#update_rounding_mode>`_
+        guide.
+
+        Raises:
+            ValueError: for invalid value types.
+
+        """
+        return self._properties.get("defaultRoundingMode")
+
+    @default_rounding_mode.setter
+    def default_rounding_mode(self, value):
+        possible_values = [
+            "ROUNDING_MODE_UNSPECIFIED",
+            "ROUND_HALF_AWAY_FROM_ZERO",
+            "ROUND_HALF_EVEN",
+        ]
+        if not isinstance(value, str) and value is not None:
+            raise ValueError("Pass a string, or None")
+        if value is None:
+            self._properties["defaultRoundingMode"] = "ROUNDING_MODE_UNSPECIFIED"
+        if value not in possible_values and value is not None:
             raise ValueError(
-                "Time Travel Window should be from 48 to 168 hours (2 to 7 days)"
+                f'rounding mode needs to be one of {",".join(possible_values)}'
             )
-        if hours % 24 != 0:
-            raise ValueError("Time Travel Window should be multiple of 24")
-        self._properties["maxTimeTravelHours"] = hours
+        if value:
+            self._properties["defaultRoundingMode"] = value
 
     @property
     def project(self):

tests/system/test_client.py

@@ -238,11 +238,6 @@ def test_create_dataset(self):
         self.assertEqual(dataset.dataset_id, DATASET_ID)
         self.assertEqual(dataset.project, Config.CLIENT.project)
 
-    def test_create_dataset_max_time_travel_hours(self):
-        DATASET_ID = _make_dataset_id("create_ci_dataset")
-        dataset = self.temp_dataset(DATASET_ID, max_time_travel_hours=24 * 2)
-        self.assertEqual(int(dataset.max_time_travel_hours), 24 * 2)
-
     def test_get_dataset(self):
         dataset_id = _make_dataset_id("get_dataset")
         client = Config.CLIENT
@@ -270,6 +265,13 @@ def test_get_dataset(self):
         self.assertEqual(got.friendly_name, "Friendly")
         self.assertEqual(got.description, "Description")
 
+    def test_create_dataset_with_default_rounding_mode(self):
+        DATASET_ID = _make_dataset_id("create_dataset_rounding_mode")
+        dataset = self.temp_dataset(DATASET_ID, default_rounding_mode="ROUND_HALF_EVEN")
+
+        self.assertTrue(_dataset_exists(dataset))
+        self.assertEqual(dataset.default_rounding_mode, "ROUND_HALF_EVEN")
+
     def test_update_dataset(self):
         dataset = self.temp_dataset(_make_dataset_id("update_dataset"))
         self.assertTrue(_dataset_exists(dataset))
@@ -2297,8 +2299,8 @@ def temp_dataset(self, dataset_id, *args, **kwargs):
         dataset = Dataset(dataset_ref)
         if kwargs.get("location"):
             dataset.location = kwargs.get("location")
-        if kwargs.get("max_time_travel_hours"):
-            dataset.max_time_travel_hours = kwargs.get("max_time_travel_hours")
+        if kwargs.get("default_rounding_mode"):
+            dataset.default_rounding_mode = kwargs.get("default_rounding_mode")
         dataset = helpers.retry_403(Config.CLIENT.create_dataset)(dataset)
         self.to_delete.append(dataset)
         return dataset

tests/unit/test_create_dataset.py

@@ -63,6 +63,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
         "datasetId": "starry-skies",
         "tableId": "northern-hemisphere",
     }
+    DEFAULT_ROUNDING_MODE = "ROUND_HALF_EVEN"
     RESOURCE = {
         "datasetReference": {"projectId": PROJECT, "datasetId": DS_ID},
         "etag": "etag",
@@ -73,6 +74,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
         "defaultTableExpirationMs": "3600",
         "labels": LABELS,
         "access": [{"role": "OWNER", "userByEmail": USER_EMAIL}, {"view": VIEW}],
+        "defaultRoundingMode": DEFAULT_ROUNDING_MODE,
     }
     conn = client._connection = make_connection(RESOURCE)
     entries = [
@@ -88,8 +90,9 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
     before.default_table_expiration_ms = 3600
     before.location = LOCATION
     before.labels = LABELS
+    before.default_rounding_mode = DEFAULT_ROUNDING_MODE
     after = client.create_dataset(before)
-
+    print(after)
     assert after.dataset_id == DS_ID
     assert after.project == PROJECT
     assert after.etag == RESOURCE["etag"]
@@ -99,6 +102,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
     assert after.location == LOCATION
     assert after.default_table_expiration_ms == 3600
     assert after.labels == LABELS
+    assert after.default_rounding_mode == DEFAULT_ROUNDING_MODE
 
     conn.api_request.assert_called_once_with(
         method="POST",
@@ -109,6 +113,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
             "friendlyName": FRIENDLY_NAME,
             "location": LOCATION,
             "defaultTableExpirationMs": "3600",
+            "defaultRoundingMode": DEFAULT_ROUNDING_MODE,
             "access": [
                 {"role": "OWNER", "userByEmail": USER_EMAIL},
                 {"view": VIEW, "role": None},
@@ -365,3 +370,106 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(PROJECT, DS_ID, LOCATION)
             mock.call(method="GET", path=get_path, timeout=DEFAULT_TIMEOUT),
         ]
     )
+
+
+def test_create_dataset_with_default_rounding_mode_if_value_is_none(
+    PROJECT, DS_ID, LOCATION
+):
+    default_rounding_mode = None
+    path = "/projects/%s/datasets" % PROJECT
+    resource = {
+        "datasetReference": {"projectId": PROJECT, "datasetId": DS_ID},
+        "etag": "etag",
+        "id": "{}:{}".format(PROJECT, DS_ID),
+        "location": LOCATION,
+    }
+    client = make_client(location=LOCATION)
+    conn = client._connection = make_connection(resource)
+
+    ds_ref = DatasetReference(PROJECT, DS_ID)
+    before = Dataset(ds_ref)
+    before.default_rounding_mode = default_rounding_mode
+    after = client.create_dataset(before)
+
+    print(40 * "-")
+    print(after)
+    print(40 * "-")
+    assert after.dataset_id == DS_ID
+    assert after.project == PROJECT
+    assert after.default_rounding_mode is None
+
+    conn.api_request.assert_called_once_with(
+        method="POST",
+        path=path,
+        data={
+            "datasetReference": {"projectId": PROJECT, "datasetId": DS_ID},
+            "labels": {},
+            "location": LOCATION,
+            "defaultRoundingMode": "ROUNDING_MODE_UNSPECIFIED",
+        },
+        timeout=DEFAULT_TIMEOUT,
+    )
+
+
+def test_create_dataset_with_default_rounding_mode_if_value_is_not_string(
+    PROJECT, DS_ID, LOCATION
+):
+    default_rounding_mode = 10
+    ds_ref = DatasetReference(PROJECT, DS_ID)
+    dataset = Dataset(ds_ref)
+    with pytest.raises(ValueError) as e:
+        dataset.default_rounding_mode = default_rounding_mode
+    assert str(e.value) == "Pass a string, or None"
+
+
+def test_create_dataset_with_default_rounding_mode_if_value_is_not_valid(
+    PROJECT, DS_ID
+):
+    default_rounding_mode = "ROUND_HALF_AWAY_FROM_ZEROS"
+    ds_ref = DatasetReference(PROJECT, DS_ID)
+    dataset = Dataset(ds_ref)
+    with pytest.raises(ValueError) as e:
+        dataset.default_rounding_mode = default_rounding_mode
+    assert (
+        str(e.value)
+        == "rounding mode needs to be one of ROUNDING_MODE_UNSPECIFIED,ROUND_HALF_AWAY_FROM_ZERO,ROUND_HALF_EVEN"
+    )
+
+
+def test_create_dataset_with_default_rounding_mode_if_value_is_valid(
+    PROJECT, DS_ID, LOCATION
+):
+    default_rounding_mode = "ROUND_HALF_AWAY_FROM_ZERO"
+    path = "/projects/%s/datasets" % PROJECT
+    resource = {
+        "datasetReference": {"projectId": PROJECT, "datasetId": DS_ID},
+        "etag": "etag",
+        "id": "{}:{}".format(PROJECT, DS_ID),
+        "location": LOCATION,
+    }
+    client = make_client(location=LOCATION)
+    conn = client._connection = make_connection(resource)
+
+    ds_ref = DatasetReference(PROJECT, DS_ID)
+    before = Dataset(ds_ref)
+    before.default_rounding_mode = default_rounding_mode
+    after = client.create_dataset(before)
+
+    print(40 * "-")
+    print(after)
+    print(40 * "-")
+    assert after.dataset_id == DS_ID
+    assert after.project == PROJECT
+    assert after.default_rounding_mode is None
+
+    conn.api_request.assert_called_once_with(
+        method="POST",
+        path=path,
+        data={
+            "datasetReference": {"projectId": PROJECT, "datasetId": DS_ID},
+            "labels": {},
+            "location": LOCATION,
+            "defaultRoundingMode": default_rounding_mode,
+        },
+        timeout=DEFAULT_TIMEOUT,
+    )