feat: add support dataset.max_time_travel_hours

Gaurang033 · Gaurang033 · commit 56f9dbee7757 · 2023-10-29T14:57:07.000-04:00
diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py
@@ -501,7 +501,9 @@ def from_api_repr(cls, resource: dict) -> "AccessEntry":
         if len(entry) != 0:
             raise ValueError("Entry has unexpected keys remaining.", entry)
 
-        return cls(role, entity_type, entity_id)
+        config = cls(role, entity_type, entity_id)
+        config._properties = copy.deepcopy(resource)
+        return config
 
 
 class Dataset(object):
@@ -525,6 +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",
     }
 
@@ -533,6 +536,27 @@ def __init__(self, dataset_ref) -> None:
             dataset_ref = DatasetReference.from_string(dataset_ref)
         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")
+
+    @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:
+            raise ValueError(
+                "Time Travel Window should be from 48 to 168 hours (2 to 7 days)"
+            )
+        if hours % 24 != 0:
+            raise ValueError("Time Travel Window should be multiple of 24")
+        self._properties["maxTimeTravelHours"] = hours
+
     @property
     def default_rounding_mode(self):
         """Union[str, None]: defaultRoundingMode of the dataset as set by the user
diff --git a/tests/system/test_client.py b/tests/system/test_client.py
@@ -238,6 +238,11 @@ 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
@@ -2299,6 +2304,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")
 
diff --git a/tests/unit/test_create_dataset.py b/tests/unit/test_create_dataset.py
@@ -466,3 +466,69 @@ def test_create_dataset_with_default_rounding_mode_if_value_is_in_possible_value
         },
         timeout=DEFAULT_TIMEOUT,
     )
+
+
+def test_create_dataset_with_max_time_travel_hours(
+    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.max_time_travel_hours = 24*3
+    after = client.create_dataset(before)
+
+    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,
+            "maxTimeTravelHours": 24*3,
+        },
+        timeout=DEFAULT_TIMEOUT,
+    )
+
+
+def test_create_dataset_with_max_time_travel_hours_not_multiple_of_24(
+    PROJECT, DS_ID, LOCATION
+):
+    ds_ref = DatasetReference(PROJECT, DS_ID)
+    dataset = Dataset(ds_ref)
+    with pytest.raises(ValueError) as e:
+        dataset.max_time_travel_hours = 50
+    assert str(e.value) == "Time Travel Window should be multiple of 24"
+
+
+def test_create_dataset_with_max_time_travel_hours_is_less_than_2_days(
+    PROJECT, DS_ID, LOCATION
+):
+    ds_ref = DatasetReference(PROJECT, DS_ID)
+    dataset = Dataset(ds_ref)
+    with pytest.raises(ValueError) as e:
+        dataset.max_time_travel_hours = 24
+    assert str(e.value) == "Time Travel Window should be from 48 to 168 hours (2 to 7 days)"
+
+
+def test_create_dataset_with_max_time_travel_hours_is_greater_than_7_days(
+    PROJECT, DS_ID, LOCATION
+):
+    ds_ref = DatasetReference(PROJECT, DS_ID)
+    dataset = Dataset(ds_ref)
+    with pytest.raises(ValueError) as e:
+        dataset.max_time_travel_hours = 24
+    assert str(e.value) == "Time Travel Window should be from 48 to 168 hours (2 to 7 days)"
