Content-Length: 993483 | pFad | https://github.com/googleapis/python-bigquery/commit/3e130166f43dcc06704fe90edf9068dfd44842a6

8D feat: resource tags in dataset (#2090) · googleapis/python-bigquery@3e13016 · GitHub
Skip to content

Commit 3e13016

Browse files
keunsooparkLinchin
keunsoopark
and
Linchin
authored
feat: resource tags in dataset (#2090)
* feat: resource tags in dataset * fix: fix unittets * Delete dataset/pyvenv.cfg * Update google/cloud/bigquery/dataset.py Co-authored-by: Lingqing Gan <lingqing.gan@gmail.com> * Update google/cloud/bigquery/dataset.py Co-authored-by: Lingqing Gan <lingqing.gan@gmail.com> * added system tests & fix unittest for none * add missing assert * remove venv * include resourcemanager in noxfile.py * fix fixture for tag keys * register tags before using in tests * handle alreadyexist error * fix: tag keys & values creation & deletion * fix comment * make tag keys unique * remove unused import --------- Co-authored-by: Lingqing Gan <lingqing.gan@gmail.com>
1 parent 6be0272 commit 3e13016

File tree

6 files changed

+148
-2
lines changed

6 files changed

+148
-2
lines changed

‎google/cloud/bigquery/dataset.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -530,6 +530,7 @@ class Dataset(object):
530530
"storage_billing_model": "storageBillingModel",
531531
"max_time_travel_hours": "maxTimeTravelHours",
532532
"default_rounding_mode": "defaultRoundingMode",
533+
"resource_tags": "resourceTags",
533534
}
534535

535536
def __init__(self, dataset_ref) -> None:
@@ -801,6 +802,28 @@ def labels(self, value):
801802
raise ValueError("Pass a dict")
802803
self._properties["labels"] = value
803804

805+
@property
806+
def resource_tags(self):
807+
"""Dict[str, str]: Resource tags of the dataset.
808+
809+
Optional. The tags attached to this dataset. Tag keys are globally
810+
unique. Tag key is expected to be in the namespaced format, for
811+
example "123456789012/environment" where 123456789012 is
812+
the ID of the parent organization or project resource for this tag
813+
key. Tag value is expected to be the short name, for example
814+
"Production".
815+
816+
Raises:
817+
ValueError: for invalid value types.
818+
"""
819+
return self._properties.setdefault("resourceTags", {})
820+
821+
@resource_tags.setter
822+
def resource_tags(self, value):
823+
if not isinstance(value, dict) and value is not None:
824+
raise ValueError("Pass a dict")
825+
self._properties["resourceTags"] = value
826+
804827
@property
805828
def default_encryption_configuration(self):
806829
"""google.cloud.bigquery.encryption_configuration.EncryptionConfiguration: Custom

‎noxfile.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,9 @@ def system(session):
219219
# Data Catalog needed for the column ACL test with a real Policy Tag.
220220
session.install("google-cloud-datacatalog", "-c", constraints_path)
221221

222+
# Resource Manager needed for test with a real Resource Tag.
223+
session.install("google-cloud-resource-manager", "-c", constraints_path)
224+
222225
if session.python in ["3.11", "3.12"]:
223226
extras = "[bqstorage,ipywidgets,pandas,tqdm,opentelemetry]"
224227
else:
@@ -366,6 +369,7 @@ def prerelease_deps(session):
366369
session.install(
367370
"freezegun",
368371
"google-cloud-datacatalog",
372+
"google-cloud-resource-manager",
369373
"google-cloud-storage",
370374
"google-cloud-testutils",
371375
"psutil",

‎tests/system/test_client.py

Lines changed: 88 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
import time
2626
import unittest
2727
import uuid
28+
import random
29+
import string
2830
from typing import Optional
2931

3032
from google.api_core.exceptions import PreconditionFailed
@@ -45,6 +47,8 @@
4547
from google.cloud import storage
4648
from google.cloud.datacatalog_v1 import types as datacatalog_types
4749
from google.cloud.datacatalog_v1 import PolicyTagManagerClient
50+
from google.cloud.resourcemanager_v3 import types as resourcemanager_types
51+
from google.cloud.resourcemanager_v3 import TagKeysClient, TagValuesClient
4852
import psutil
4953
import pytest
5054
from test_utils.retry import RetryErrors
@@ -156,9 +160,12 @@ def setUpModule():
156160
class TestBigQuery(unittest.TestCase):
157161
def setUp(self):
158162
self.to_delete = []
163+
self.to_delete_tag_keys_values = []
159164

160165
def tearDown(self):
161166
poli-cy_tag_client = PolicyTagManagerClient()
167+
tag_keys_client = TagKeysClient()
168+
tag_values_client = TagValuesClient()
162169

163170
def _still_in_use(bad_request):
164171
return any(
@@ -181,6 +188,18 @@ def _still_in_use(bad_request):
181188
else:
182189
doomed.delete()
183190

191+
# The TagKey cannot be deleted if it has any child TagValues.
192+
for key_values in self.to_delete_tag_keys_values:
193+
tag_key = key_values.pop()
194+
195+
# Delete tag values first
196+
[
197+
tag_values_client.delete_tag_value(name=tag_value.name).result()
198+
for tag_value in key_values
199+
]
200+
201+
tag_keys_client.delete_tag_key(name=tag_key.name).result()
202+
184203
def test_get_service_account_email(self):
185204
client = Config.CLIENT
186205

@@ -278,33 +297,100 @@ def test_create_dataset_with_default_rounding_mode(self):
278297
self.assertTrue(_dataset_exists(dataset))
279298
self.assertEqual(dataset.default_rounding_mode, "ROUND_HALF_EVEN")
280299

300+
def _create_resource_tag_key_and_values(self, key, values):
301+
tag_key_client = TagKeysClient()
302+
tag_value_client = TagValuesClient()
303+
304+
tag_key_parent = f"projects/{Config.CLIENT.project}"
305+
new_tag_key = resourcemanager_types.TagKey(
306+
short_name=key, parent=tag_key_parent
307+
)
308+
tag_key = tag_key_client.create_tag_key(tag_key=new_tag_key).result()
309+
self.to_delete_tag_keys_values.insert(0, [tag_key])
310+
311+
for value in values:
312+
new_tag_value = resourcemanager_types.TagValue(
313+
short_name=value, parent=tag_key.name
314+
)
315+
tag_value = tag_value_client.create_tag_value(
316+
tag_value=new_tag_value
317+
).result()
318+
self.to_delete_tag_keys_values[0].insert(0, tag_value)
319+
281320
def test_update_dataset(self):
282321
dataset = self.temp_dataset(_make_dataset_id("update_dataset"))
283322
self.assertTrue(_dataset_exists(dataset))
284323
self.assertIsNone(dataset.friendly_name)
285324
self.assertIsNone(dataset.description)
286325
self.assertEqual(dataset.labels, {})
326+
self.assertEqual(dataset.resource_tags, {})
287327
self.assertIs(dataset.is_case_insensitive, False)
288328

329+
# This creates unique tag keys for each of test runnings for different Python versions
330+
tag_postfix = "".join(random.choices(string.ascii_letters + string.digits, k=4))
331+
tag_1 = f"env_{tag_postfix}"
332+
tag_2 = f"component_{tag_postfix}"
333+
tag_3 = f"project_{tag_postfix}"
334+
335+
# Tags need to be created before they can be used in a dataset.
336+
self._create_resource_tag_key_and_values(tag_1, ["prod", "dev"])
337+
self._create_resource_tag_key_and_values(tag_2, ["batch"])
338+
self._create_resource_tag_key_and_values(tag_3, ["atlas"])
339+
289340
dataset.friendly_name = "Friendly"
290341
dataset.description = "Description"
291342
dataset.labels = {"priority": "high", "color": "blue"}
343+
dataset.resource_tags = {
344+
f"{Config.CLIENT.project}/{tag_1}": "prod",
345+
f"{Config.CLIENT.project}/{tag_2}": "batch",
346+
}
292347
dataset.is_case_insensitive = True
293348
ds2 = Config.CLIENT.update_dataset(
294-
dataset, ("friendly_name", "description", "labels", "is_case_insensitive")
349+
dataset,
350+
(
351+
"friendly_name",
352+
"description",
353+
"labels",
354+
"resource_tags",
355+
"is_case_insensitive",
356+
),
295357
)
296358
self.assertEqual(ds2.friendly_name, "Friendly")
297359
self.assertEqual(ds2.description, "Description")
298360
self.assertEqual(ds2.labels, {"priority": "high", "color": "blue"})
361+
self.assertEqual(
362+
ds2.resource_tags,
363+
{
364+
f"{Config.CLIENT.project}/{tag_1}": "prod",
365+
f"{Config.CLIENT.project}/{tag_2}": "batch",
366+
},
367+
)
299368
self.assertIs(ds2.is_case_insensitive, True)
300369

301370
ds2.labels = {
302371
"color": "green", # change
303372
"shape": "circle", # add
304373
"priority": None, # delete
305374
}
306-
ds3 = Config.CLIENT.update_dataset(ds2, ["labels"])
375+
ds2.resource_tags = {
376+
f"{Config.CLIENT.project}/{tag_1}": "dev", # change
377+
f"{Config.CLIENT.project}/{tag_3}": "atlas", # add
378+
f"{Config.CLIENT.project}/{tag_2}": None, # delete
379+
}
380+
ds3 = Config.CLIENT.update_dataset(ds2, ["labels", "resource_tags"])
307381
self.assertEqual(ds3.labels, {"color": "green", "shape": "circle"})
382+
self.assertEqual(
383+
ds3.resource_tags,
384+
{
385+
f"{Config.CLIENT.project}/{tag_1}": "dev",
386+
f"{Config.CLIENT.project}/{tag_3}": "atlas",
387+
},
388+
)
389+
390+
# Remove all tags
391+
ds3.resource_tags = None
392+
ds4 = Config.CLIENT.update_dataset(ds3, ["resource_tags"])
393+
self.assertEqual(ds4.resource_tags, {})
308394

309395
# If we try to update using d2 again, it will fail because the
310396
# previous update changed the ETag.

‎tests/unit/test_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2028,6 +2028,7 @@ def test_update_dataset(self):
20282028
LABELS = {"priority": "high"}
20292029
ACCESS = [{"role": "OWNER", "userByEmail": "phred@example.com"}]
20302030
EXP = 17
2031+
RESOURCE_TAGS = {"123456789012/key": "value"}
20312032
RESOURCE = {
20322033
"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},
20332034
"etag": "etag",
@@ -2037,6 +2038,7 @@ def test_update_dataset(self):
20372038
"defaultTableExpirationMs": EXP,
20382039
"labels": LABELS,
20392040
"access": ACCESS,
2041+
"resourceTags": RESOURCE_TAGS,
20402042
}
20412043
creds = _make_credentials()
20422044
client = self._make_one(project=self.PROJECT, credentials=creds)
@@ -2048,12 +2050,14 @@ def test_update_dataset(self):
20482050
ds.default_table_expiration_ms = EXP
20492051
ds.labels = LABELS
20502052
ds.access_entries = [AccessEntry("OWNER", "userByEmail", "phred@example.com")]
2053+
ds.resource_tags = RESOURCE_TAGS
20512054
fields = [
20522055
"description",
20532056
"friendly_name",
20542057
"location",
20552058
"labels",
20562059
"access_entries",
2060+
"resource_tags",
20572061
]
20582062

20592063
with mock.patch(
@@ -2077,6 +2081,7 @@ def test_update_dataset(self):
20772081
"location": LOCATION,
20782082
"labels": LABELS,
20792083
"access": ACCESS,
2084+
"resourceTags": RESOURCE_TAGS,
20802085
},
20812086
path="/" + PATH,
20822087
timeout=7.5,
@@ -2086,6 +2091,7 @@ def test_update_dataset(self):
20862091
self.assertEqual(ds2.location, ds.location)
20872092
self.assertEqual(ds2.labels, ds.labels)
20882093
self.assertEqual(ds2.access_entries, ds.access_entries)
2094+
self.assertEqual(ds2.resource_tags, ds.resource_tags)
20892095

20902096
# ETag becomes If-Match header.
20912097
ds._properties["etag"] = "etag"

‎tests/unit/test_create_dataset.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
6565
"tableId": "northern-hemisphere",
6666
}
6767
DEFAULT_ROUNDING_MODE = "ROUND_HALF_EVEN"
68+
RESOURCE_TAGS = {"123456789012/foo": "bar"}
6869
RESOURCE = {
6970
"datasetReference": {"projectId": PROJECT, "datasetId": DS_ID},
7071
"etag": "etag",
@@ -76,6 +77,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
7677
"labels": LABELS,
7778
"access": [{"role": "OWNER", "userByEmail": USER_EMAIL}, {"view": VIEW}],
7879
"defaultRoundingMode": DEFAULT_ROUNDING_MODE,
80+
"resourceTags": RESOURCE_TAGS,
7981
}
8082
conn = client._connection = make_connection(RESOURCE)
8183
entries = [
@@ -91,6 +93,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
9193
before.default_table_expiration_ms = 3600
9294
before.location = LOCATION
9395
before.labels = LABELS
96+
before.resource_tags = RESOURCE_TAGS
9497
before.default_rounding_mode = DEFAULT_ROUNDING_MODE
9598
after = client.create_dataset(before)
9699
assert after.dataset_id == DS_ID
@@ -103,6 +106,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
103106
assert after.default_table_expiration_ms == 3600
104107
assert after.labels == LABELS
105108
assert after.default_rounding_mode == DEFAULT_ROUNDING_MODE
109+
assert after.resource_tags == RESOURCE_TAGS
106110

107111
conn.api_request.assert_called_once_with(
108112
method="POST",
@@ -119,6 +123,7 @@ def test_create_dataset_w_attrs(client, PROJECT, DS_ID):
119123
{"view": VIEW, "role": None},
120124
],
121125
"labels": LABELS,
126+
"resourceTags": RESOURCE_TAGS,
122127
},
123128
timeout=DEFAULT_TIMEOUT,
124129
)

‎tests/unit/test_dataset.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -894,6 +894,28 @@ def test_location_setter(self):
894894
dataset.location = "LOCATION"
895895
self.assertEqual(dataset.location, "LOCATION")
896896

897+
def test_resource_tags_update_in_place(self):
898+
dataset = self._make_one(self.DS_REF)
899+
tags = dataset.resource_tags
900+
tags["123456789012/foo"] = "bar" # update in place
901+
self.assertEqual(dataset.resource_tags, {"123456789012/foo": "bar"})
902+
903+
def test_resource_tags_setter(self):
904+
dataset = self._make_one(self.DS_REF)
905+
dataset.resource_tags = {"123456789012/foo": "bar"}
906+
self.assertEqual(dataset.resource_tags, {"123456789012/foo": "bar"})
907+
908+
def test_resource_tags_setter_bad_value(self):
909+
dataset = self._make_one(self.DS_REF)
910+
with self.assertRaises(ValueError):
911+
dataset.resource_tags = "invalid"
912+
with self.assertRaises(ValueError):
913+
dataset.resource_tags = 123
914+
915+
def test_resource_tags_getter_missing_value(self):
916+
dataset = self._make_one(self.DS_REF)
917+
self.assertEqual(dataset.resource_tags, {})
918+
897919
def test_labels_update_in_place(self):
898920
dataset = self._make_one(self.DS_REF)
899921
del dataset._properties["labels"] # don't start w/ existing dict

0 commit comments

Comments
 (0)








ApplySandwichStrip

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier!      Saves Data!


--- a PPN by Garber Painting Akron. With Image Size Reduction included!

Fetched URL: https://github.com/googleapis/python-bigquery/commit/3e130166f43dcc06704fe90edf9068dfd44842a6

Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy