feat: add support for dataset.default_rounding_mode

Gaurang Shah · Gaurang Shah · commit 590d65e4f12e · 2023-10-18T14:31:20.000-04:00
diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py
@@ -527,13 +527,46 @@ class Dataset(object):
         "friendly_name": "friendlyName",
         "default_encryption_configuration": "defaultEncryptionConfiguration",
         "storage_billing_model": "storageBillingModel",
+        "default_rounding_mode": "defaultRoundingMode",
     }
 
     def __init__(self, dataset_ref) -> None:
         if isinstance(dataset_ref, str):
             dataset_ref = DatasetReference.from_string(dataset_ref)
         self._properties = {"datasetReference": dataset_ref.to_api_repr(), "labels": {}}
 
+    @property
+    def default_rounding_mode(self):
+        """
+        optional[str]: it could have one of the following possible value.
+        ROUNDING_MODE_UNSPECIFIED: Unspecified will default to using ROUND_HALF_AWAY_FROM_ZERO.
+        ROUND_HALF_AWAY_FROM_ZERO: ROUND_HALF_AWAY_FROM_ZERO rounds half values away from zero when applying precision
+                                    and scale upon writing of NUMERIC and BIGNUMERIC values.
+                                    For Scale: 0 1.1, 1.2, 1.3, 1.4 => 1 1.5, 1.6, 1.7, 1.8, 1.9 => 2
+        ROUND_HALF_EVEN: ROUND_HALF_EVEN rounds half values to the nearest even value when applying precision and scale
+                        upon writing of NUMERIC and BIGNUMERIC values.
+                        For Scale: 0 1.1, 1.2, 1.3, 1.4 => 1 1.5 => 2 1.6, 1.7, 1.8, 1.9 => 2 2.5 => 2
+        """
+        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 not in possible_values:
+            raise ValueError(
+                f'rounding mode needs to be one of {",".join(possible_values)}'
+            )
+        if value:
+            self._properties["defaultRoundingMode"] = value
+        else:
+            self._properties["defaultRoundingMode"] = "ROUNDING_MODE_UNSPECIFIED"
+
     @property
     def project(self):
         """str: Project ID of the project bound to the dataset."""
diff --git a/tests/system/test_client.py b/tests/system/test_client.py
@@ -265,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))
@@ -2286,12 +2293,17 @@ def test_nested_table_to_arrow(self):
         self.assertTrue(pyarrow.types.is_list(record_col[1].type))
         self.assertTrue(pyarrow.types.is_int64(record_col[1].type.value_type))
 
-    def temp_dataset(self, dataset_id, location=None):
+    def temp_dataset(self, dataset_id, *args, **kwargs):
         project = Config.CLIENT.project
         dataset_ref = bigquery.DatasetReference(project, dataset_id)
         dataset = Dataset(dataset_ref)
-        if location:
-            dataset.location = location
+        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
