Skip to content

Commit 18bb443

Browse files
plamutplamut
authored
feat: enable unsetting policy tags on schema fields (#703)
* feat: enable unsetting policy tags on schema fields * Adjust API representation for STRUCT schema fields * De-dup logic for converting None policy tags
1 parent 0b20015 commit 18bb443

File tree

6 files changed

+279
-32
lines changed

6 files changed

+279
-32
lines changed

‎google/cloud/bigquery/schema.py

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
"""Schemas for BigQuery tables / queries."""
1616

1717
import collections
18+
from typing import Optional
1819

1920
from google.cloud.bigquery_v2 import types
2021

@@ -105,7 +106,26 @@ def __init__(
105106
if max_length is not _DEFAULT_VALUE:
106107
self._properties["maxLength"] = max_length
107108
self._fields = tuple(fields)
108-
self._policy_tags = policy_tags
109+
110+
self._policy_tags = self._determine_policy_tags(field_type, policy_tags)
111+
112+
@staticmethod
113+
def _determine_policy_tags(
114+
field_type: str, given_policy_tags: Optional["PolicyTagList"]
115+
) -> Optional["PolicyTagList"]:
116+
"""Return the given policy tags, or their suitable representation if `None`.
117+
118+
Args:
119+
field_type: The type of the schema field.
120+
given_policy_tags: The policy tags to maybe ajdust.
121+
"""
122+
if given_policy_tags is not None:
123+
return given_policy_tags
124+
125+
if field_type is not None and field_type.upper() in _STRUCT_TYPES:
126+
return None
127+
128+
return PolicyTagList()
109129

110130
@staticmethod
111131
def __get_int(api_repr, name):
@@ -126,18 +146,24 @@ def from_api_repr(cls, api_repr: dict) -> "SchemaField":
126146
Returns:
127147
google.cloud.biquery.schema.SchemaField: The ``SchemaField`` object.
128148
"""
149+
field_type = api_repr["type"].upper()
150+
129151
# Handle optional properties with default values
130152
mode = api_repr.get("mode", "NULLABLE")
131153
description = api_repr.get("description", _DEFAULT_VALUE)
132154
fields = api_repr.get("fields", ())
133155

156+
policy_tags = cls._determine_policy_tags(
157+
field_type, PolicyTagList.from_api_repr(api_repr.get("policyTags"))
158+
)
159+
134160
return cls(
135-
field_type=api_repr["type"].upper(),
161+
field_type=field_type,
136162
fields=[cls.from_api_repr(f) for f in fields],
137163
mode=mode.upper(),
138164
description=description,
139165
name=api_repr["name"],
140-
policy_tags=PolicyTagList.from_api_repr(api_repr.get("policyTags")),
166+
policy_tags=policy_tags,
141167
precision=cls.__get_int(api_repr, "precision"),
142168
scale=cls.__get_int(api_repr, "scale"),
143169
max_length=cls.__get_int(api_repr, "maxLength"),
@@ -218,9 +244,9 @@ def to_api_repr(self) -> dict:
218244
# add this to the serialized representation.
219245
if self.field_type.upper() in _STRUCT_TYPES:
220246
answer["fields"] = [f.to_api_repr() for f in self.fields]
221-
222-
# If this contains a policy tag definition, include that as well:
223-
if self.policy_tags is not None:
247+
else:
248+
# Explicitly include policy tag definition (we must not do it for RECORD
249+
# fields, because those are not leaf fields).
224250
answer["policyTags"] = self.policy_tags.to_api_repr()
225251

226252
# Done; return the serialized dictionary.
@@ -244,14 +270,19 @@ def _key(self):
244270
field_type = f"{field_type}({self.precision}, {self.scale})"
245271
else:
246272
field_type = f"{field_type}({self.precision})"
273+
274+
policy_tags = (
275+
() if self._policy_tags is None else tuple(sorted(self._policy_tags.names))
276+
)
277+
247278
return (
248279
self.name,
249280
field_type,
250281
# Mode is always str, if not given it defaults to a str value
251282
self.mode.upper(), # pytype: disable=attribute-error
252283
self.description,
253284
self._fields,
254-
self._policy_tags,
285+
policy_tags,
255286
)
256287

257288
def to_standard_sql(self) -> types.StandardSqlField:

‎tests/system/test_client.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -653,6 +653,56 @@ def test_update_table_schema(self):
653653
self.assertEqual(found.field_type, expected.field_type)
654654
self.assertEqual(found.mode, expected.mode)
655655

656+
def test_unset_table_schema_attributes(self):
657+
from google.cloud.bigquery.schema import PolicyTagList
658+
659+
dataset = self.temp_dataset(_make_dataset_id("unset_policy_tags"))
660+
table_id = "test_table"
661+
policy_tags = PolicyTagList(
662+
names=[
663+
"projects/{}/locations/us/taxonomies/1/policyTags/2".format(
664+
Config.CLIENT.project
665+
),
666+
]
667+
)
668+
669+
schema = [
670+
bigquery.SchemaField("full_name", "STRING", mode="REQUIRED"),
671+
bigquery.SchemaField(
672+
"secret_int",
673+
"INTEGER",
674+
mode="REQUIRED",
675+
description="This field is numeric",
676+
policy_tags=policy_tags,
677+
),
678+
]
679+
table_arg = Table(dataset.table(table_id), schema=schema)
680+
self.assertFalse(_table_exists(table_arg))
681+
682+
table = helpers.retry_403(Config.CLIENT.create_table)(table_arg)
683+
self.to_delete.insert(0, table)
684+
685+
self.assertTrue(_table_exists(table))
686+
self.assertEqual(policy_tags, table.schema[1].policy_tags)
687+
688+
# Amend the schema to replace the policy tags
689+
new_schema = table.schema[:]
690+
old_field = table.schema[1]
691+
new_schema[1] = bigquery.SchemaField(
692+
name=old_field.name,
693+
field_type=old_field.field_type,
694+
mode=old_field.mode,
695+
description=None,
696+
fields=old_field.fields,
697+
policy_tags=None,
698+
)
699+
700+
table.schema = new_schema
701+
updated_table = Config.CLIENT.update_table(table, ["schema"])
702+
703+
self.assertFalse(updated_table.schema[1].description) # Empty string or None.
704+
self.assertEqual(updated_table.schema[1].policy_tags.names, ())
705+
656706
def test_update_table_clustering_configuration(self):
657707
dataset = self.temp_dataset(_make_dataset_id("update_table"))
658708

‎tests/unit/job/test_load_config.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -434,11 +434,13 @@ def test_schema_setter_fields(self):
434434
"name": "full_name",
435435
"type": "STRING",
436436
"mode": "REQUIRED",
437+
"policyTags": {"names": []},
437438
}
438439
age_repr = {
439440
"name": "age",
440441
"type": "INTEGER",
441442
"mode": "REQUIRED",
443+
"policyTags": {"names": []},
442444
}
443445
self.assertEqual(
444446
config._properties["load"]["schema"], {"fields": [full_name_repr, age_repr]}
@@ -451,11 +453,13 @@ def test_schema_setter_valid_mappings_list(self):
451453
"name": "full_name",
452454
"type": "STRING",
453455
"mode": "REQUIRED",
456+
"policyTags": {"names": []},
454457
}
455458
age_repr = {
456459
"name": "age",
457460
"type": "INTEGER",
458461
"mode": "REQUIRED",
462+
"policyTags": {"names": []},
459463
}
460464
schema = [full_name_repr, age_repr]
461465
config.schema = schema

‎tests/unit/test_client.py

Lines changed: 43 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1019,8 +1019,18 @@ def test_create_table_w_schema_and_query(self):
10191019
{
10201020
"schema": {
10211021
"fields": [
1022-
{"name": "full_name", "type": "STRING", "mode": "REQUIRED"},
1023-
{"name": "age", "type": "INTEGER", "mode": "REQUIRED"},
1022+
{
1023+
"name": "full_name",
1024+
"type": "STRING",
1025+
"mode": "REQUIRED",
1026+
"policyTags": {"names": []},
1027+
},
1028+
{
1029+
"name": "age",
1030+
"type": "INTEGER",
1031+
"mode": "REQUIRED",
1032+
"policyTags": {"names": []},
1033+
},
10241034
]
10251035
},
10261036
"view": {"query": query},
@@ -1054,8 +1064,18 @@ def test_create_table_w_schema_and_query(self):
10541064
},
10551065
"schema": {
10561066
"fields": [
1057-
{"name": "full_name", "type": "STRING", "mode": "REQUIRED"},
1058-
{"name": "age", "type": "INTEGER", "mode": "REQUIRED"},
1067+
{
1068+
"name": "full_name",
1069+
"type": "STRING",
1070+
"mode": "REQUIRED",
1071+
"policyTags": {"names": []},
1072+
},
1073+
{
1074+
"name": "age",
1075+
"type": "INTEGER",
1076+
"mode": "REQUIRED",
1077+
"policyTags": {"names": []},
1078+
},
10591079
]
10601080
},
10611081
"view": {"query": query, "useLegacySql": False},
@@ -2000,12 +2020,14 @@ def test_update_table(self):
20002020
"type": "STRING",
20012021
"mode": "REQUIRED",
20022022
"description": None,
2023+
"policyTags": {"names": []},
20032024
},
20042025
{
20052026
"name": "age",
20062027
"type": "INTEGER",
20072028
"mode": "REQUIRED",
20082029
"description": "New field description",
2030+
"policyTags": {"names": []},
20092031
},
20102032
]
20112033
},
@@ -2047,12 +2069,14 @@ def test_update_table(self):
20472069
"type": "STRING",
20482070
"mode": "REQUIRED",
20492071
"description": None,
2072+
"policyTags": {"names": []},
20502073
},
20512074
{
20522075
"name": "age",
20532076
"type": "INTEGER",
20542077
"mode": "REQUIRED",
20552078
"description": "New field description",
2079+
"policyTags": {"names": []},
20562080
},
20572081
]
20582082
},
@@ -2173,14 +2197,21 @@ def test_update_table_w_query(self):
21732197
"type": "STRING",
21742198
"mode": "REQUIRED",
21752199
"description": None,
2200+
"policyTags": {"names": []},
21762201
},
21772202
{
21782203
"name": "age",
21792204
"type": "INTEGER",
21802205
"mode": "REQUIRED",
21812206
"description": "this is a column",
2207+
"policyTags": {"names": []},
2208+
},
2209+
{
2210+
"name": "country",
2211+
"type": "STRING",
2212+
"mode": "NULLABLE",
2213+
"policyTags": {"names": []},
21822214
},
2183-
{"name": "country", "type": "STRING", "mode": "NULLABLE"},
21842215
]
21852216
}
21862217
schema = [
@@ -6516,10 +6547,10 @@ def test_load_table_from_dataframe(self):
65166547
assert field["type"] == table_field.field_type
65176548
assert field["mode"] == table_field.mode
65186549
assert len(field.get("fields", [])) == len(table_field.fields)
6550+
assert field["policyTags"]["names"] == []
65196551
# Omit unnecessary fields when they come from getting the table
65206552
# (not passed in via job_config)
65216553
assert "description" not in field
6522-
assert "policyTags" not in field
65236554

65246555
@unittest.skipIf(pandas is None, "Requires `pandas`")
65256556
@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
@@ -7718,18 +7749,21 @@ def test_schema_to_json_with_file_path(self):
77187749
"description": "quarter",
77197750
"mode": "REQUIRED",
77207751
"name": "qtr",
7752+
"policyTags": {"names": []},
77217753
"type": "STRING",
77227754
},
77237755
{
77247756
"description": "sales representative",
77257757
"mode": "NULLABLE",
77267758
"name": "rep",
7759+
"policyTags": {"names": []},
77277760
"type": "STRING",
77287761
},
77297762
{
77307763
"description": "total sales",
77317764
"mode": "NULLABLE",
77327765
"name": "sales",
7766+
"policyTags": {"names": []},
77337767
"type": "FLOAT",
77347768
},
77357769
]
@@ -7762,18 +7796,21 @@ def test_schema_to_json_with_file_object(self):
77627796
"description": "quarter",
77637797
"mode": "REQUIRED",
77647798
"name": "qtr",
7799+
"policyTags": {"names": []},
77657800
"type": "STRING",
77667801
},
77677802
{
77687803
"description": "sales representative",
77697804
"mode": "NULLABLE",
77707805
"name": "rep",
7806+
"policyTags": {"names": []},
77717807
"type": "STRING",
77727808
},
77737809
{
77747810
"description": "total sales",
77757811
"mode": "NULLABLE",
77767812
"name": "sales",
7813+
"policyTags": {"names": []},
77777814
"type": "FLOAT",
77787815
},
77797816
]

‎tests/unit/test_external_config.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,14 @@ def test_to_api_repr_base(self):
7878
ec.schema = [schema.SchemaField("full_name", "STRING", mode="REQUIRED")]
7979

8080
exp_schema = {
81-
"fields": [{"name": "full_name", "type": "STRING", "mode": "REQUIRED"}]
81+
"fields": [
82+
{
83+
"name": "full_name",
84+
"type": "STRING",
85+
"mode": "REQUIRED",
86+
"policyTags": {"names": []},
87+
}
88+
]
8289
}
8390
got_resource = ec.to_api_repr()
8491
exp_resource = {

0 commit comments

Comments
 (0)