From e7839096e7b2483f146092aa770d907ca6b341af Mon Sep 17 00:00:00 2001 From: hrkh Date: Sat, 21 Dec 2024 13:17:01 +0900 Subject: [PATCH 01/10] feat: support resource_tags for table --- google/cloud/bigquery/table.py | 20 ++++++++++++++++++++ tests/system/test_client.py | 16 +++++++++++++++- tests/unit/test_client.py | 8 ++++++-- tests/unit/test_table.py | 27 +++++++++++++++++++++++++++ 4 files changed, 68 insertions(+), 3 deletions(-) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index 80ab330ba..b01c37cc7 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -408,6 +408,7 @@ class Table(_TableBase): "require_partition_filter": "requirePartitionFilter", "table_constraints": "tableConstraints", "max_staleness": "maxStaleness", + "resource_tags": "resourceTags", } def __init__(self, table_ref, schema=None) -> None: @@ -1023,6 +1024,25 @@ def table_constraints(self) -> Optional["TableConstraints"]: table_constraints = TableConstraints.from_api_repr(table_constraints) return table_constraints + @property + def resource_tags(self): + """Dict[str, str]: Resource tags for the table. + + See: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table.FIELDS.resource_tags + + Raises: + ValueError: For invalid value types. + """ + return self._properties.setdefault( + self._PROPERTY_TO_API_FIELD["resource_tags"], {} + ) + + @resource_tags.setter + def resource_tags(self, value): + if not isinstance(value, dict): + raise ValueError("resource_tags must be a dict") + self._properties[self._PROPERTY_TO_API_FIELD["resource_tags"]] = value + @classmethod def from_string(cls, full_table_id: str) -> "Table": """Construct a table from fully-qualified table ID. diff --git a/tests/system/test_client.py b/tests/system/test_client.py index 95c679a14..3a36faebd 100644 --- a/tests/system/test_client.py +++ b/tests/system/test_client.py @@ -658,14 +658,19 @@ def test_update_table(self): table.friendly_name = "Friendly" table.description = "Description" table.labels = {"priority": "high", "color": "blue"} + table.resource_tags = {"123456789012/owner": "Alice", "123456789012/env": "dev"} table2 = Config.CLIENT.update_table( - table, ["friendly_name", "description", "labels"] + table, ["friendly_name", "description", "labels", "resource_tags"] ) self.assertEqual(table2.friendly_name, "Friendly") self.assertEqual(table2.description, "Description") self.assertEqual(table2.labels, {"priority": "high", "color": "blue"}) + self.assertEqual( + table2.resource_tags, + {"123456789012/owner": "Alice", "123456789012/env": "dev"}, + ) table2.description = None table2.labels = { @@ -673,9 +678,18 @@ def test_update_table(self): "shape": "circle", # add "priority": None, # delete } + table2.resource_tags = { + "123456789012/owner": "Bob", # change + "123456789012/classification": "public", # add + "123456789012/env": None, # delete + } table3 = Config.CLIENT.update_table(table2, ["description", "labels"]) self.assertIsNone(table3.description) self.assertEqual(table3.labels, {"color": "green", "shape": "circle"}) + self.assertEqual( + table3.resource_tags, + {"123456789012/owner": "Bob", "123456789012/classification": "public"}, + ) # If we try to update using table2 again, it will fail because the # previous update changed the ETag. diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index cd336b73f..22a40527e 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -2314,6 +2314,7 @@ def test_update_table(self): "description": description, "friendlyName": title, "labels": {"x": "y"}, + "resourceTags": {"123456789012/key": "value"}, } ) schema = [ @@ -2337,7 +2338,8 @@ def test_update_table(self): table.description = description table.friendly_name = title table.labels = {"x": "y"} - fields = ["schema", "description", "friendly_name", "labels"] + table.resource_tags = {"123456789012/key": "value"} + fields = ["schema", "description", "friendly_name", "labels", "resource_tags"] with mock.patch( "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes" ) as final_attributes: @@ -2369,14 +2371,16 @@ def test_update_table(self): "description": description, "friendlyName": title, "labels": {"x": "y"}, + "resourceTags": {"123456789012/key": "value"}, } conn.api_request.assert_called_once_with( - method="PATCH", data=sent, path="/" + path, timeout=7.5 + method="PATCH", path="/" + path, data=sent, timeout=7.5 ) self.assertEqual(updated_table.description, table.description) self.assertEqual(updated_table.friendly_name, table.friendly_name) self.assertEqual(updated_table.schema, table.schema) self.assertEqual(updated_table.labels, table.labels) + self.assertEqual(updated_table.resource_tags, table.resource_tags) # ETag becomes If-Match header. table._properties["etag"] = "etag" diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 3824da226..5eb2f0a68 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -1458,6 +1458,33 @@ def test_encryption_configuration_setter(self): table.encryption_configuration = None self.assertIsNone(table.encryption_configuration) + def test_resource_tags_getter_empty(self): + dataset = DatasetReference(self.PROJECT, self.DS_ID) + table_ref = dataset.table(self.TABLE_NAME) + table = self._make_one(table_ref) + self.assertEqual(table.resource_tags, {}) + + def test_resource_tags_update_in_place(self): + dataset = DatasetReference(self.PROJECT, self.DS_ID) + table_ref = dataset.table(self.TABLE_NAME) + table = self._make_one(table_ref) + table.resource_tags["123456789012/key"] = "value" + self.assertEqual(table.resource_tags, {"123456789012/key": "value"}) + + def test_resource_tags_setter(self): + dataset = DatasetReference(self.PROJECT, self.DS_ID) + table_ref = dataset.table(self.TABLE_NAME) + table = self._make_one(table_ref) + table.resource_tags = {"123456789012/key": "value"} + self.assertEqual(table.resource_tags, {"123456789012/key": "value"}) + + def test_resource_tags_setter_bad_value(self): + dataset = DatasetReference(self.PROJECT, self.DS_ID) + table_ref = dataset.table(self.TABLE_NAME) + table = self._make_one(table_ref) + with self.assertRaises(ValueError): + table.resource_tags = None + def test___repr__(self): from google.cloud.bigquery.table import TableReference From a194256ad2b2407507681d2474bd03fad8da0a9a Mon Sep 17 00:00:00 2001 From: hrkh Date: Sat, 21 Dec 2024 17:47:22 +0900 Subject: [PATCH 02/10] fix: system test for resource tags --- google/cloud/bigquery/table.py | 4 +-- noxfile.py | 5 ++- tests/system/test_client.py | 62 ++++++++++++++++++++++++++++++---- 3 files changed, 61 insertions(+), 10 deletions(-) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index b01c37cc7..93e645120 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -1039,8 +1039,8 @@ def resource_tags(self): @resource_tags.setter def resource_tags(self, value): - if not isinstance(value, dict): - raise ValueError("resource_tags must be a dict") + if not isinstance(value, dict) and value is not None: + raise ValueError("resource_tags must be a dict or None") self._properties[self._PROPERTY_TO_API_FIELD["resource_tags"]] = value @classmethod diff --git a/noxfile.py b/noxfile.py index 750a6b459..bf0dad371 100644 --- a/noxfile.py +++ b/noxfile.py @@ -219,6 +219,9 @@ def system(session): # Data Catalog needed for the column ACL test with a real Policy Tag. session.install("google-cloud-datacatalog", "-c", constraints_path) + # Resource Manager needed for test with a real Resource Tag. + session.install("google-cloud-resource-manager", "-c", constraints_path) + if session.python in ["3.11", "3.12"]: extras = "[bqstorage,ipywidgets,pandas,tqdm,opentelemetry]" else: @@ -366,7 +369,7 @@ def prerelease_deps(session): session.install( "freezegun", "google-cloud-datacatalog", - "google-cloud-storage", + "google-cloud-resource-manager" "google-cloud-storage", "google-cloud-testutils", "psutil", "pytest", diff --git a/tests/system/test_client.py b/tests/system/test_client.py index 3a36faebd..a63d31dc3 100644 --- a/tests/system/test_client.py +++ b/tests/system/test_client.py @@ -45,6 +45,8 @@ from google.cloud import storage from google.cloud.datacatalog_v1 import types as datacatalog_types from google.cloud.datacatalog_v1 import PolicyTagManagerClient +from google.cloud.resourcemanager_v3 import types as resourcemanager_types +from google.cloud.resourcemanager_v3 import TagKeysClient, TagValuesClient import psutil import pytest from test_utils.retry import RetryErrors @@ -159,6 +161,8 @@ def setUp(self): def tearDown(self): policy_tag_client = PolicyTagManagerClient() + tag_keys_client = TagKeysClient() + tag_values_client = TagValuesClient() def _still_in_use(bad_request): return any( @@ -178,6 +182,10 @@ def _still_in_use(bad_request): retry_in_use(Config.CLIENT.delete_table)(doomed) elif isinstance(doomed, datacatalog_types.Taxonomy): policy_tag_client.delete_taxonomy(name=doomed.name) + elif isinstance(doomed, resourcemanager_types.TagKey): + tag_keys_client.delete_tag_key(name=doomed.name).result() + elif isinstance(doomed, resourcemanager_types.TagValue): + tag_values_client.delete_tag_value(name=doomed.name).result() else: doomed.delete() @@ -646,6 +654,30 @@ def test_list_tables(self): def test_update_table(self): dataset = self.temp_dataset(_make_dataset_id("update_table")) + def _create_resource_tag_key_and_values(key, values): + tag_key_client = TagKeysClient() + tag_value_client = TagValuesClient() + + tag_key_parent = f"projects/{Config.CLIENT.project}" + new_tag_key = resourcemanager_types.TagKey( + short_name=key, parent=tag_key_parent + ) + tag_key = tag_key_client.create_tag_key(tag_key=new_tag_key).result() + self.to_delete.insert(0, tag_key) + + for value in values: + new_tag_value = resourcemanager_types.TagValue( + short_name=value, parent=tag_key.name + ) + tag_value = tag_value_client.create_tag_value( + tag_value=new_tag_value + ).result() + self.to_delete.insert(0, tag_value) + + _create_resource_tag_key_and_values("owner", ["Alice", "Bob"]) + _create_resource_tag_key_and_values("classification", ["public"]) + _create_resource_tag_key_and_values("env", ["dev"]) + TABLE_NAME = "test_table" table_arg = Table(dataset.table(TABLE_NAME), schema=SCHEMA) self.assertFalse(_table_exists(table_arg)) @@ -658,7 +690,10 @@ def test_update_table(self): table.friendly_name = "Friendly" table.description = "Description" table.labels = {"priority": "high", "color": "blue"} - table.resource_tags = {"123456789012/owner": "Alice", "123456789012/env": "dev"} + table.resource_tags = { + f"{Config.CLIENT.project}/owner": "Alice", + f"{Config.CLIENT.project}/env": "dev", + } table2 = Config.CLIENT.update_table( table, ["friendly_name", "description", "labels", "resource_tags"] @@ -669,7 +704,10 @@ def test_update_table(self): self.assertEqual(table2.labels, {"priority": "high", "color": "blue"}) self.assertEqual( table2.resource_tags, - {"123456789012/owner": "Alice", "123456789012/env": "dev"}, + { + f"{Config.CLIENT.project}/owner": "Alice", + f"{Config.CLIENT.project}/env": "dev", + }, ) table2.description = None @@ -679,18 +717,28 @@ def test_update_table(self): "priority": None, # delete } table2.resource_tags = { - "123456789012/owner": "Bob", # change - "123456789012/classification": "public", # add - "123456789012/env": None, # delete + f"{Config.CLIENT.project}/owner": "Bob", # change + f"{Config.CLIENT.project}/classification": "public", # add + f"{Config.CLIENT.project}/env": None, # delete } - table3 = Config.CLIENT.update_table(table2, ["description", "labels"]) + table3 = Config.CLIENT.update_table( + table2, ["description", "labels", "resource_tags"] + ) self.assertIsNone(table3.description) self.assertEqual(table3.labels, {"color": "green", "shape": "circle"}) self.assertEqual( table3.resource_tags, - {"123456789012/owner": "Bob", "123456789012/classification": "public"}, + { + f"{Config.CLIENT.project}/owner": "Bob", + f"{Config.CLIENT.project}/classification": "public", + }, ) + # Delete resource tag bindings. + table3.resource_tags = None + table4 = Config.CLIENT.update_table(table3, ["resource_tags"]) + self.assertEqual(table4.resource_tags, {}) + # If we try to update using table2 again, it will fail because the # previous update changed the ETag. table2.description = "no good" From 75fde55b8e7ab95b3916c0768f1eab32feb8f477 Mon Sep 17 00:00:00 2001 From: hrkh Date: Sat, 21 Dec 2024 17:49:38 +0900 Subject: [PATCH 03/10] fix: typo --- noxfile.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index bf0dad371..e08956b11 100644 --- a/noxfile.py +++ b/noxfile.py @@ -369,7 +369,8 @@ def prerelease_deps(session): session.install( "freezegun", "google-cloud-datacatalog", - "google-cloud-resource-manager" "google-cloud-storage", + "google-cloud-resource-manager", + "google-cloud-storage", "google-cloud-testutils", "psutil", "pytest", From 873caf80736031e35ecfce996993051f5acf042c Mon Sep 17 00:00:00 2001 From: hrkh Date: Sat, 21 Dec 2024 19:25:44 +0900 Subject: [PATCH 04/10] fix: unit test --- tests/unit/test_table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py index 5eb2f0a68..17996892b 100644 --- a/tests/unit/test_table.py +++ b/tests/unit/test_table.py @@ -1483,7 +1483,7 @@ def test_resource_tags_setter_bad_value(self): table_ref = dataset.table(self.TABLE_NAME) table = self._make_one(table_ref) with self.assertRaises(ValueError): - table.resource_tags = None + table.resource_tags = 12345 def test___repr__(self): from google.cloud.bigquery.table import TableReference From 43cc9969acfc3ca9b4effa533991d5c194b3dc5b Mon Sep 17 00:00:00 2001 From: Lingqing Gan Date: Tue, 7 Jan 2025 07:13:24 +0800 Subject: [PATCH 05/10] Update tests/unit/test_client.py --- tests/unit/test_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 22a40527e..9377c2692 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -2374,7 +2374,7 @@ def test_update_table(self): "resourceTags": {"123456789012/key": "value"}, } conn.api_request.assert_called_once_with( - method="PATCH", path="/" + path, data=sent, timeout=7.5 + method="PATCH", data=sent, path="/" + path, timeout=7.5 ) self.assertEqual(updated_table.description, table.description) self.assertEqual(updated_table.friendly_name, table.friendly_name) From 19f66d5e4daa8f19046c37af6f76adc67dc03868 Mon Sep 17 00:00:00 2001 From: Lingqing Gan Date: Tue, 7 Jan 2025 07:14:43 +0800 Subject: [PATCH 06/10] Update google/cloud/bigquery/table.py --- google/cloud/bigquery/table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index 93e645120..fd0390416 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -1033,7 +1033,7 @@ def resource_tags(self): Raises: ValueError: For invalid value types. """ - return self._properties.setdefault( + return self._properties.get( self._PROPERTY_TO_API_FIELD["resource_tags"], {} ) From e42b4cf01cbc0bdabee9ac057e662432408a7dd5 Mon Sep 17 00:00:00 2001 From: Lingqing Gan Date: Tue, 7 Jan 2025 07:14:56 +0800 Subject: [PATCH 07/10] Update google/cloud/bigquery/table.py --- google/cloud/bigquery/table.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index fd0390416..5ca2e6bda 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -1029,9 +1029,6 @@ def resource_tags(self): """Dict[str, str]: Resource tags for the table. See: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table.FIELDS.resource_tags - - Raises: - ValueError: For invalid value types. """ return self._properties.get( self._PROPERTY_TO_API_FIELD["resource_tags"], {} From aaa125cbfeeb24cbfb855588897b9d110d343dd1 Mon Sep 17 00:00:00 2001 From: Lingqing Gan Date: Thu, 16 Jan 2025 12:55:27 -0800 Subject: [PATCH 08/10] Update google/cloud/bigquery/table.py --- google/cloud/bigquery/table.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index 5ca2e6bda..59885a23b 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -1030,7 +1030,7 @@ def resource_tags(self): See: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table.FIELDS.resource_tags """ - return self._properties.get( + return self._properties.setdefault( self._PROPERTY_TO_API_FIELD["resource_tags"], {} ) From 6eb99f2e6a562dc21ae30d79018f7070dff08b0f Mon Sep 17 00:00:00 2001 From: hrkh Date: Sat, 18 Jan 2025 16:35:20 +0900 Subject: [PATCH 09/10] fix: append random string suffix to resource tags to prevent test conflicts --- tests/system/test_client.py | 52 ++++++++++++------------------------- 1 file changed, 17 insertions(+), 35 deletions(-) diff --git a/tests/system/test_client.py b/tests/system/test_client.py index 4c04db1f5..30e9f94a3 100644 --- a/tests/system/test_client.py +++ b/tests/system/test_client.py @@ -185,10 +185,6 @@ def _still_in_use(bad_request): retry_in_use(Config.CLIENT.delete_table)(doomed) elif isinstance(doomed, datacatalog_types.Taxonomy): policy_tag_client.delete_taxonomy(name=doomed.name) - elif isinstance(doomed, resourcemanager_types.TagKey): - tag_keys_client.delete_tag_key(name=doomed.name).result() - elif isinstance(doomed, resourcemanager_types.TagValue): - tag_values_client.delete_tag_value(name=doomed.name).result() else: doomed.delete() @@ -736,29 +732,15 @@ def test_list_tables(self): def test_update_table(self): dataset = self.temp_dataset(_make_dataset_id("update_table")) - def _create_resource_tag_key_and_values(key, values): - tag_key_client = TagKeysClient() - tag_value_client = TagValuesClient() - - tag_key_parent = f"projects/{Config.CLIENT.project}" - new_tag_key = resourcemanager_types.TagKey( - short_name=key, parent=tag_key_parent - ) - tag_key = tag_key_client.create_tag_key(tag_key=new_tag_key).result() - self.to_delete.insert(0, tag_key) - - for value in values: - new_tag_value = resourcemanager_types.TagValue( - short_name=value, parent=tag_key.name - ) - tag_value = tag_value_client.create_tag_value( - tag_value=new_tag_value - ).result() - self.to_delete.insert(0, tag_value) + # This creates unique tag keys for each of test runnings for different Python versions + tag_postfix = "".join(random.choices(string.ascii_letters + string.digits, k=4)) + tag_1 = f"owner_{tag_postfix}" + tag_2 = f"classification_{tag_postfix}" + tag_3 = f"env_{tag_postfix}" - _create_resource_tag_key_and_values("owner", ["Alice", "Bob"]) - _create_resource_tag_key_and_values("classification", ["public"]) - _create_resource_tag_key_and_values("env", ["dev"]) + self._create_resource_tag_key_and_values(tag_1, ["Alice", "Bob"]) + self._create_resource_tag_key_and_values(tag_2, ["public"]) + self._create_resource_tag_key_and_values(tag_3, ["dev"]) TABLE_NAME = "test_table" table_arg = Table(dataset.table(TABLE_NAME), schema=SCHEMA) @@ -773,8 +755,8 @@ def _create_resource_tag_key_and_values(key, values): table.description = "Description" table.labels = {"priority": "high", "color": "blue"} table.resource_tags = { - f"{Config.CLIENT.project}/owner": "Alice", - f"{Config.CLIENT.project}/env": "dev", + f"{Config.CLIENT.project}/{tag_1}": "Alice", + f"{Config.CLIENT.project}/{tag_3}": "dev", } table2 = Config.CLIENT.update_table( @@ -787,8 +769,8 @@ def _create_resource_tag_key_and_values(key, values): self.assertEqual( table2.resource_tags, { - f"{Config.CLIENT.project}/owner": "Alice", - f"{Config.CLIENT.project}/env": "dev", + f"{Config.CLIENT.project}/{tag_1}": "Alice", + f"{Config.CLIENT.project}/{tag_3}": "dev", }, ) @@ -799,9 +781,9 @@ def _create_resource_tag_key_and_values(key, values): "priority": None, # delete } table2.resource_tags = { - f"{Config.CLIENT.project}/owner": "Bob", # change - f"{Config.CLIENT.project}/classification": "public", # add - f"{Config.CLIENT.project}/env": None, # delete + f"{Config.CLIENT.project}/{tag_1}": "Bob", # change + f"{Config.CLIENT.project}/{tag_2}": "public", # add + f"{Config.CLIENT.project}/{tag_3}": None, # delete } table3 = Config.CLIENT.update_table( table2, ["description", "labels", "resource_tags"] @@ -811,8 +793,8 @@ def _create_resource_tag_key_and_values(key, values): self.assertEqual( table3.resource_tags, { - f"{Config.CLIENT.project}/owner": "Bob", - f"{Config.CLIENT.project}/classification": "public", + f"{Config.CLIENT.project}/{tag_1}": "Bob", + f"{Config.CLIENT.project}/{tag_2}": "public", }, ) From 0fc85d6567e621d6a4a7ba18d9552db1b4e87e43 Mon Sep 17 00:00:00 2001 From: Lingqing Gan Date: Tue, 21 Jan 2025 11:11:38 -0800 Subject: [PATCH 10/10] Update google/cloud/bigquery/table.py --- google/cloud/bigquery/table.py | 1 + 1 file changed, 1 insertion(+) diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py index 75a8bf6bd..934a28cfc 100644 --- a/google/cloud/bigquery/table.py +++ b/google/cloud/bigquery/table.py @@ -1042,6 +1042,7 @@ def resource_tags(self, value): raise ValueError("resource_tags must be a dict or None") self._properties[self._PROPERTY_TO_API_FIELD["resource_tags"]] = value + @property def external_catalog_table_options( self, ) -> Optional[external_config.ExternalCatalogTableOptions]: