Skip to content

Commit

Permalink
feat: add support for dataset.default_rounding_mode
Browse files Browse the repository at this point in the history
  • Loading branch information
Gaurang033 committed Oct 20, 2023
1 parent 535c982 commit 114955e
Show file tree
Hide file tree
Showing 3 changed files with 152 additions and 25 deletions.
51 changes: 34 additions & 17 deletions google/cloud/bigquery/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -527,7 +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",
}

def __init__(self, dataset_ref) -> None:
Expand All @@ -536,25 +536,42 @@ def __init__(self, dataset_ref) -> None:
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")
def default_rounding_mode(self):
"""Union[str, None]: defaultRoundingMode of the dataset as set by the user
(defaults to :data:`None`).
Set the value to one of ``'ROUND_HALF_AWAY_FROM_ZERO'``, ``'ROUND_HALF_EVEN'``, or
``'ROUNDING_MODE_UNSPECIFIED'``.
@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:
See `default rounding mode
<https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#Dataset.FIELDS.default_rounding_mode>`_
in REST API docs and `updating the default rounding model
<https://cloud.google.com/bigquery/docs/updating-datasets#update_rounding_mode>`_
guide.
Raises:
ValueError: for invalid value types.
"""
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 is None:
self._properties["defaultRoundingMode"] = "ROUNDING_MODE_UNSPECIFIED"
if value not in possible_values and value is not None:
raise ValueError(
"Time Travel Window should be from 48 to 168 hours (2 to 7 days)"
f'rounding mode needs to be one of {",".join(possible_values)}'
)
if hours % 24 != 0:
raise ValueError("Time Travel Window should be multiple of 24")
self._properties["maxTimeTravelHours"] = hours
if value:
self._properties["defaultRoundingMode"] = value

@property
def project(self):
Expand Down
16 changes: 9 additions & 7 deletions tests/system/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,11 +238,6 @@ 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 @@ -270,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))
Expand Down Expand Up @@ -2297,8 +2299,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")
dataset = helpers.retry_403(Config.CLIENT.create_dataset)(dataset)
self.to_delete.append(dataset)
return dataset
Expand Down
110 changes: 109 additions & 1 deletion tests/unit/test_create_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
"datasetId": "starry-skies",
"tableId": "northern-hemisphere",
}
DEFAULT_ROUNDING_MODE = "ROUND_HALF_EVEN"
RESOURCE = {
"datasetReference": {"projectId": PROJECT, "datasetId": DS_ID},
"etag": "etag",
Expand All @@ -73,6 +74,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
"defaultTableExpirationMs": "3600",
"labels": LABELS,
"access": [{"role": "OWNER", "userByEmail": USER_EMAIL}, {"view": VIEW}],
"defaultRoundingMode": DEFAULT_ROUNDING_MODE,
}
conn = client._connection = make_connection(RESOURCE)
entries = [
Expand All @@ -88,8 +90,9 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
before.default_table_expiration_ms = 3600
before.location = LOCATION
before.labels = LABELS
before.default_rounding_mode = DEFAULT_ROUNDING_MODE
after = client.create_dataset(before)

print(after)
assert after.dataset_id == DS_ID
assert after.project == PROJECT
assert after.etag == RESOURCE["etag"]
Expand All @@ -99,6 +102,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
assert after.location == LOCATION
assert after.default_table_expiration_ms == 3600
assert after.labels == LABELS
assert after.default_rounding_mode == DEFAULT_ROUNDING_MODE

conn.api_request.assert_called_once_with(
method="POST",
Expand All @@ -109,6 +113,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
"friendlyName": FRIENDLY_NAME,
"location": LOCATION,
"defaultTableExpirationMs": "3600",
"defaultRoundingMode": DEFAULT_ROUNDING_MODE,
"access": [
{"role": "OWNER", "userByEmail": USER_EMAIL},
{"view": VIEW, "role": None},
Expand Down Expand Up @@ -365,3 +370,106 @@ def test_create_dataset_alreadyexists_w_exists_ok_true(PROJECT, DS_ID, LOCATION)
mock.call(method="GET", path=get_path, timeout=DEFAULT_TIMEOUT),
]
)


def test_create_dataset_with_default_rounding_mode_if_value_is_none(
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.default_rounding_mode = default_rounding_mode
after = client.create_dataset(before)

print(40 * "-")
print(after)
print(40 * "-")
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,
"defaultRoundingMode": "ROUNDING_MODE_UNSPECIFIED",
},
timeout=DEFAULT_TIMEOUT,
)


def test_create_dataset_with_default_rounding_mode_if_value_is_not_string(
PROJECT, DS_ID, LOCATION
):
default_rounding_mode = 10
ds_ref = DatasetReference(PROJECT, DS_ID)
dataset = Dataset(ds_ref)
with pytest.raises(ValueError) as e:
dataset.default_rounding_mode = default_rounding_mode
assert str(e.value) == "Pass a string, or None"


def test_create_dataset_with_default_rounding_mode_if_value_is_not_valid(
PROJECT, DS_ID
):
default_rounding_mode = "ROUND_HALF_AWAY_FROM_ZEROS"
ds_ref = DatasetReference(PROJECT, DS_ID)
dataset = Dataset(ds_ref)
with pytest.raises(ValueError) as e:
dataset.default_rounding_mode = default_rounding_mode
assert (
str(e.value)
== "rounding mode needs to be one of ROUNDING_MODE_UNSPECIFIED,ROUND_HALF_AWAY_FROM_ZERO,ROUND_HALF_EVEN"
)


def test_create_dataset_with_default_rounding_mode_if_value_is_valid(
PROJECT, DS_ID, LOCATION
):
default_rounding_mode = "ROUND_HALF_AWAY_FROM_ZERO"
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.default_rounding_mode = default_rounding_mode
after = client.create_dataset(before)

print(40 * "-")
print(after)
print(40 * "-")
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,
"defaultRoundingMode": default_rounding_mode,
},
timeout=DEFAULT_TIMEOUT,
)

0 comments on commit 114955e

Please sign in to comment.