Skip to content

Commit

Permalink
feat: add support dataset.max_time_travel_hours (#1683)
Browse files Browse the repository at this point in the history
* feat: add support dataset.max_time_travel_hours

* Update tests/unit/test_create_dataset.py

* Update tests/unit/test_create_dataset.py

* Update google/cloud/bigquery/dataset.py

* update test_create_dataset_with_max_time_travel_hours

---------

Co-authored-by: Lingqing Gan <lingqing.gan@gmail.com>
  • Loading branch information
Gaurang033 and Linchin committed Oct 31, 2023
1 parent 49bfd12 commit f22eff2
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 0 deletions.
23 changes: 23 additions & 0 deletions google/cloud/bigquery/dataset.py
Expand Up @@ -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",
}

Expand All @@ -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
Expand Down
7 changes: 7 additions & 0 deletions tests/system/test_client.py
Expand Up @@ -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
Expand Down Expand Up @@ -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")

Expand Down
79 changes: 79 additions & 0 deletions tests/unit/test_create_dataset.py
Expand Up @@ -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"

0 comments on commit f22eff2

Please sign in to comment.