Skip to content

feat: support resource_tags for table #2093

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 13 commits into from
Jan 21, 2025
Merged
17 changes: 17 additions & 0 deletions google/cloud/bigquery/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -409,6 +409,7 @@ class Table(_TableBase):
"require_partition_filter": "requirePartitionFilter",
"table_constraints": "tableConstraints",
"max_staleness": "maxStaleness",
"resource_tags": "resourceTags",
"external_catalog_table_options": "externalCatalogTableOptions",
}

Expand Down Expand Up @@ -1025,6 +1026,22 @@ 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
"""
return self._properties.setdefault(
self._PROPERTY_TO_API_FIELD["resource_tags"], {}
)

@resource_tags.setter
def resource_tags(self, value):
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

@property
def external_catalog_table_options(
self,
Expand Down
44 changes: 42 additions & 2 deletions tests/system/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -732,6 +732,16 @@ def test_list_tables(self):
def test_update_table(self):
dataset = self.temp_dataset(_make_dataset_id("update_table"))

# 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}"

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)
self.assertFalse(_table_exists(table_arg))
Expand All @@ -744,24 +754,54 @@ def test_update_table(self):
table.friendly_name = "Friendly"
table.description = "Description"
table.labels = {"priority": "high", "color": "blue"}
table.resource_tags = {
f"{Config.CLIENT.project}/{tag_1}": "Alice",
f"{Config.CLIENT.project}/{tag_3}": "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,
{
f"{Config.CLIENT.project}/{tag_1}": "Alice",
f"{Config.CLIENT.project}/{tag_3}": "dev",
},
)

table2.description = None
table2.labels = {
"color": "green", # change
"shape": "circle", # add
"priority": None, # delete
}
table3 = Config.CLIENT.update_table(table2, ["description", "labels"])
table2.resource_tags = {
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"]
)
self.assertIsNone(table3.description)
self.assertEqual(table3.labels, {"color": "green", "shape": "circle"})
self.assertEqual(
table3.resource_tags,
{
f"{Config.CLIENT.project}/{tag_1}": "Bob",
f"{Config.CLIENT.project}/{tag_2}": "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.
Expand Down
6 changes: 5 additions & 1 deletion tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -2320,6 +2320,7 @@ def test_update_table(self):
"description": description,
"friendlyName": title,
"labels": {"x": "y"},
"resourceTags": {"123456789012/key": "value"},
}
)
schema = [
Expand All @@ -2343,7 +2344,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:
Expand Down Expand Up @@ -2375,6 +2377,7 @@ 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
Expand All @@ -2383,6 +2386,7 @@ def test_update_table(self):
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"
Expand Down
27 changes: 27 additions & 0 deletions tests/unit/test_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -1481,6 +1481,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 = 12345

def test___repr__(self):
from google.cloud.bigquery.table import TableReference

Expand Down