diff --git a/google/cloud/bigquery/dataset.py b/google/cloud/bigquery/dataset.py index af94784a4..726a2a17a 100644 --- a/google/cloud/bigquery/dataset.py +++ b/google/cloud/bigquery/dataset.py @@ -525,6 +525,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 +534,28 @@ 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), and 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(f"max_time_travel_hours must be an integer. Got {hours}") + 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 index 09606590e..9660d5fa7 100644 --- 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 index 3b2e644d9..8374e6e0a 100644 --- a/tests/unit/test_create_dataset.py +++ b/tests/unit/test_create_dataset.py @@ -466,3 +466,82 @@ 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): + path = "/projects/%s/datasets" % PROJECT + max_time_travel_hours = 24 * 3 + + resource = { + "datasetReference": {"projectId": PROJECT, "datasetId": DS_ID}, + "etag": "etag", + "id": "{}:{}".format(PROJECT, DS_ID), + "location": LOCATION, + "maxTimeTravelHours": max_time_travel_hours, + } + 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 = max_time_travel_hours + after = client.create_dataset(before) + assert after.dataset_id == DS_ID + assert after.project == PROJECT + assert after.max_time_travel_hours == max_time_travel_hours + + conn.api_request.assert_called_once_with( + method="POST", + path=path, + data={ + "datasetReference": {"projectId": PROJECT, "datasetId": DS_ID}, + "labels": {}, + "location": LOCATION, + "maxTimeTravelHours": max_time_travel_hours, + }, + 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 = 192 + 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_not_int(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) == "max_time_travel_hours must be an integer. Got 50"