Skip to content

Commit

Permalink
feat: add support dataset.max_time_travel_hours
Browse files Browse the repository at this point in the history
  • Loading branch information
Gaurang033 authored and Gaurang Shah committed Oct 18, 2023
1 parent 5ceed05 commit 4241cc1
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 3 deletions.
18 changes: 18 additions & 0 deletions google/cloud/bigquery/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -527,13 +527,31 @@ class Dataset(object):
"friendly_name": "friendlyName",
"default_encryption_configuration": "defaultEncryptionConfiguration",
"storage_billing_model": "storageBillingModel",
"max_time_travel_hours": "maxTimeTravelHours"
}

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 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).
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)")

self._properties["maxTimeTravelHours"] = hours

@property
def project(self):
"""str: Project ID of the project bound to the dataset."""
Expand Down
13 changes: 10 additions & 3 deletions tests/system/test_client.py
Original file line number Diff line number Diff line change
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 @@ -2286,12 +2291,14 @@ 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")
dataset = helpers.retry_403(Config.CLIENT.create_dataset)(dataset)
self.to_delete.append(dataset)
return dataset
Expand Down

0 comments on commit 4241cc1

Please sign in to comment.