Skip to content

Commit

Permalink
Support resource_tags for Table
Browse files Browse the repository at this point in the history
  • Loading branch information
hrkh committed Dec 21, 2024
1 parent d4d39ac commit 732426d
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 3 deletions.
20 changes: 20 additions & 0 deletions google/cloud/bigquery/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down
16 changes: 15 additions & 1 deletion tests/system/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -658,24 +658,38 @@ 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 = {
"color": "green", # change
"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.
Expand Down
8 changes: 6 additions & 2 deletions tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -2314,6 +2314,7 @@ def test_update_table(self):
"description": description,
"friendlyName": title,
"labels": {"x": "y"},
"resourceTags": {"123456789012/key": "value"},
}
)
schema = [
Expand All @@ -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:
Expand Down Expand Up @@ -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"
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 @@ -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

Expand Down

0 comments on commit 732426d

Please sign in to comment.