Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 38a5c01

Browse filesBrowse files
shollymanplamut
shollyman
and
plamut
authored
feat: add support for policy tags (#77)
* feat: add support for policy tags in schema * blacken * add more unit coverage * more test cleanup * more tests * formatting * more testing of names setter * address reviewer comments * docstrings migrate from unions -> optional * stashing changes * revision to list-based representation, update tests * changes to equality and testing, towards satisfying coverage * cleanup * return copy * address api repr feedback * make PolicyTagList fully immutable * update docstring * simplify to_api_repr * remove stale doc comments Co-authored-by: Peter Lamut <plamut@users.noreply.github.com>
1 parent 23a173b commit 38a5c01
Copy full SHA for 38a5c01

File tree

Expand file treeCollapse file tree

3 files changed

+278
-5
lines changed
Open diff view settings
Filter options
Expand file treeCollapse file tree

3 files changed

+278
-5
lines changed
Open diff view settings
Collapse file

‎google/cloud/bigquery/schema.py‎

Copy file name to clipboardExpand all lines: google/cloud/bigquery/schema.py
+116-2Lines changed: 116 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,14 +62,26 @@ class SchemaField(object):
6262
6363
fields (Tuple[google.cloud.bigquery.schema.SchemaField]):
6464
subfields (requires ``field_type`` of 'RECORD').
65+
66+
policy_tags (Optional[PolicyTagList]): The policy tag list for the field.
67+
6568
"""
6669

67-
def __init__(self, name, field_type, mode="NULLABLE", description=None, fields=()):
70+
def __init__(
71+
self,
72+
name,
73+
field_type,
74+
mode="NULLABLE",
75+
description=None,
76+
fields=(),
77+
policy_tags=None,
78+
):
6879
self._name = name
6980
self._field_type = field_type
7081
self._mode = mode
7182
self._description = description
7283
self._fields = tuple(fields)
84+
self._policy_tags = policy_tags
7385

7486
@classmethod
7587
def from_api_repr(cls, api_repr):
@@ -87,12 +99,14 @@ def from_api_repr(cls, api_repr):
8799
mode = api_repr.get("mode", "NULLABLE")
88100
description = api_repr.get("description")
89101
fields = api_repr.get("fields", ())
102+
90103
return cls(
91104
field_type=api_repr["type"].upper(),
92105
fields=[cls.from_api_repr(f) for f in fields],
93106
mode=mode.upper(),
94107
description=description,
95108
name=api_repr["name"],
109+
policy_tags=PolicyTagList.from_api_repr(api_repr.get("policyTags")),
96110
)
97111

98112
@property
@@ -136,6 +150,13 @@ def fields(self):
136150
"""
137151
return self._fields
138152

153+
@property
154+
def policy_tags(self):
155+
"""Optional[google.cloud.bigquery.schema.PolicyTagList]: Policy tag list
156+
definition for this field.
157+
"""
158+
return self._policy_tags
159+
139160
def to_api_repr(self):
140161
"""Return a dictionary representing this schema field.
141162
@@ -155,6 +176,10 @@ def to_api_repr(self):
155176
if self.field_type.upper() in _STRUCT_TYPES:
156177
answer["fields"] = [f.to_api_repr() for f in self.fields]
157178

179+
# If this contains a policy tag definition, include that as well:
180+
if self.policy_tags is not None:
181+
answer["policyTags"] = self.policy_tags.to_api_repr()
182+
158183
# Done; return the serialized dictionary.
159184
return answer
160185

@@ -172,6 +197,7 @@ def _key(self):
172197
self._mode.upper(),
173198
self._description,
174199
self._fields,
200+
self._policy_tags,
175201
)
176202

177203
def to_standard_sql(self):
@@ -244,7 +270,10 @@ def _parse_schema_resource(info):
244270
mode = r_field.get("mode", "NULLABLE")
245271
description = r_field.get("description")
246272
sub_fields = _parse_schema_resource(r_field)
247-
schema.append(SchemaField(name, field_type, mode, description, sub_fields))
273+
policy_tags = PolicyTagList.from_api_repr(r_field.get("policyTags"))
274+
schema.append(
275+
SchemaField(name, field_type, mode, description, sub_fields, policy_tags)
276+
)
248277
return schema
249278

250279

@@ -291,3 +320,88 @@ def _to_schema_fields(schema):
291320
field if isinstance(field, SchemaField) else SchemaField.from_api_repr(field)
292321
for field in schema
293322
]
323+
324+
325+
class PolicyTagList(object):
326+
"""Define Policy Tags for a column.
327+
328+
Args:
329+
names (
330+
Optional[Tuple[str]]): list of policy tags to associate with
331+
the column. Policy tag identifiers are of the form
332+
`projects/*/locations/*/taxonomies/*/policyTags/*`.
333+
"""
334+
335+
def __init__(self, names=()):
336+
self._properties = {}
337+
self._properties["names"] = tuple(names)
338+
339+
@property
340+
def names(self):
341+
"""Tuple[str]: Policy tags associated with this definition.
342+
"""
343+
return self._properties.get("names", ())
344+
345+
def _key(self):
346+
"""A tuple key that uniquely describes this PolicyTagList.
347+
348+
Used to compute this instance's hashcode and evaluate equality.
349+
350+
Returns:
351+
Tuple: The contents of this :class:`~google.cloud.bigquery.schema.PolicyTagList`.
352+
"""
353+
return tuple(sorted(self._properties.items()))
354+
355+
def __eq__(self, other):
356+
if not isinstance(other, PolicyTagList):
357+
return NotImplemented
358+
return self._key() == other._key()
359+
360+
def __ne__(self, other):
361+
return not self == other
362+
363+
def __hash__(self):
364+
return hash(self._key())
365+
366+
def __repr__(self):
367+
return "PolicyTagList{}".format(self._key())
368+
369+
@classmethod
370+
def from_api_repr(cls, api_repr):
371+
"""Return a :class:`PolicyTagList` object deserialized from a dict.
372+
373+
This method creates a new ``PolicyTagList`` instance that points to
374+
the ``api_repr`` parameter as its internal properties dict. This means
375+
that when a ``PolicyTagList`` instance is stored as a property of
376+
another object, any changes made at the higher level will also appear
377+
here.
378+
379+
Args:
380+
api_repr (Mapping[str, str]):
381+
The serialized representation of the PolicyTagList, such as
382+
what is output by :meth:`to_api_repr`.
383+
384+
Returns:
385+
Optional[google.cloud.bigquery.schema.PolicyTagList]:
386+
The ``PolicyTagList`` object or None.
387+
"""
388+
if api_repr is None:
389+
return None
390+
names = api_repr.get("names", ())
391+
return cls(names=names)
392+
393+
def to_api_repr(self):
394+
"""Return a dictionary representing this object.
395+
396+
This method returns the properties dict of the ``PolicyTagList``
397+
instance rather than making a copy. This means that when a
398+
``PolicyTagList`` instance is stored as a property of another
399+
object, any changes made at the higher level will also appear here.
400+
401+
Returns:
402+
dict:
403+
A dictionary representing the PolicyTagList object in
404+
serialized form.
405+
"""
406+
answer = {"names": [name for name in self.names]}
407+
return answer
Collapse file

‎tests/system.py‎

Copy file name to clipboardExpand all lines: tests/system.py
+51Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,57 @@ def test_create_table(self):
339339
self.assertTrue(_table_exists(table))
340340
self.assertEqual(table.table_id, table_id)
341341

342+
def test_create_table_with_policy(self):
343+
from google.cloud.bigquery.schema import PolicyTagList
344+
345+
dataset = self.temp_dataset(_make_dataset_id("create_table_with_policy"))
346+
table_id = "test_table"
347+
policy_1 = PolicyTagList(
348+
names=[
349+
"projects/{}/locations/us/taxonomies/1/policyTags/2".format(
350+
Config.CLIENT.project
351+
),
352+
]
353+
)
354+
policy_2 = PolicyTagList(
355+
names=[
356+
"projects/{}/locations/us/taxonomies/3/policyTags/4".format(
357+
Config.CLIENT.project
358+
),
359+
]
360+
)
361+
362+
schema = [
363+
bigquery.SchemaField("full_name", "STRING", mode="REQUIRED"),
364+
bigquery.SchemaField(
365+
"secret_int", "INTEGER", mode="REQUIRED", policy_tags=policy_1
366+
),
367+
]
368+
table_arg = Table(dataset.table(table_id), schema=schema)
369+
self.assertFalse(_table_exists(table_arg))
370+
371+
table = retry_403(Config.CLIENT.create_table)(table_arg)
372+
self.to_delete.insert(0, table)
373+
374+
self.assertTrue(_table_exists(table))
375+
self.assertEqual(policy_1, table.schema[1].policy_tags)
376+
377+
# Amend the schema to replace the policy tags
378+
new_schema = table.schema[:]
379+
old_field = table.schema[1]
380+
new_schema[1] = bigquery.SchemaField(
381+
name=old_field.name,
382+
field_type=old_field.field_type,
383+
mode=old_field.mode,
384+
description=old_field.description,
385+
fields=old_field.fields,
386+
policy_tags=policy_2,
387+
)
388+
389+
table.schema = new_schema
390+
table2 = Config.CLIENT.update_table(table, ["schema"])
391+
self.assertEqual(policy_2, table2.schema[1].policy_tags)
392+
342393
def test_create_table_w_time_partitioning_w_clustering_fields(self):
343394
from google.cloud.bigquery.table import TimePartitioning
344395
from google.cloud.bigquery.table import TimePartitioningType
Collapse file

‎tests/unit/test_schema.py‎

Copy file name to clipboardExpand all lines: tests/unit/test_schema.py
+111-3Lines changed: 111 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,11 +63,38 @@ def test_constructor_subfields(self):
6363
self.assertIs(field._fields[0], sub_field1)
6464
self.assertIs(field._fields[1], sub_field2)
6565

66+
def test_constructor_with_policy_tags(self):
67+
from google.cloud.bigquery.schema import PolicyTagList
68+
69+
policy = PolicyTagList(names=("foo", "bar"))
70+
field = self._make_one(
71+
"test", "STRING", mode="REQUIRED", description="Testing", policy_tags=policy
72+
)
73+
self.assertEqual(field._name, "test")
74+
self.assertEqual(field._field_type, "STRING")
75+
self.assertEqual(field._mode, "REQUIRED")
76+
self.assertEqual(field._description, "Testing")
77+
self.assertEqual(field._fields, ())
78+
self.assertEqual(field._policy_tags, policy)
79+
6680
def test_to_api_repr(self):
67-
field = self._make_one("foo", "INTEGER", "NULLABLE")
81+
from google.cloud.bigquery.schema import PolicyTagList
82+
83+
policy = PolicyTagList(names=("foo", "bar"))
84+
self.assertEqual(
85+
policy.to_api_repr(), {"names": ["foo", "bar"]},
86+
)
87+
88+
field = self._make_one("foo", "INTEGER", "NULLABLE", policy_tags=policy)
6889
self.assertEqual(
6990
field.to_api_repr(),
70-
{"mode": "NULLABLE", "name": "foo", "type": "INTEGER", "description": None},
91+
{
92+
"mode": "NULLABLE",
93+
"name": "foo",
94+
"type": "INTEGER",
95+
"description": None,
96+
"policyTags": {"names": ["foo", "bar"]},
97+
},
7198
)
7299

73100
def test_to_api_repr_with_subfield(self):
@@ -111,6 +138,23 @@ def test_from_api_repr(self):
111138
self.assertEqual(field.fields[0].field_type, "INTEGER")
112139
self.assertEqual(field.fields[0].mode, "NULLABLE")
113140

141+
def test_from_api_repr_policy(self):
142+
field = self._get_target_class().from_api_repr(
143+
{
144+
"fields": [{"mode": "nullable", "name": "bar", "type": "integer"}],
145+
"name": "foo",
146+
"type": "record",
147+
"policyTags": {"names": ["one", "two"]},
148+
}
149+
)
150+
self.assertEqual(field.name, "foo")
151+
self.assertEqual(field.field_type, "RECORD")
152+
self.assertEqual(field.policy_tags.names, ("one", "two"))
153+
self.assertEqual(len(field.fields), 1)
154+
self.assertEqual(field.fields[0].name, "bar")
155+
self.assertEqual(field.fields[0].field_type, "INTEGER")
156+
self.assertEqual(field.fields[0].mode, "NULLABLE")
157+
114158
def test_from_api_repr_defaults(self):
115159
field = self._get_target_class().from_api_repr(
116160
{"name": "foo", "type": "record"}
@@ -408,7 +452,7 @@ def test___hash__not_equals(self):
408452

409453
def test___repr__(self):
410454
field1 = self._make_one("field1", "STRING")
411-
expected = "SchemaField('field1', 'STRING', 'NULLABLE', None, ())"
455+
expected = "SchemaField('field1', 'STRING', 'NULLABLE', None, (), None)"
412456
self.assertEqual(repr(field1), expected)
413457

414458

@@ -632,3 +676,67 @@ def test_valid_mapping_representation(self):
632676

633677
result = self._call_fut(schema)
634678
self.assertEqual(result, expected_schema)
679+
680+
681+
class TestPolicyTags(unittest.TestCase):
682+
@staticmethod
683+
def _get_target_class():
684+
from google.cloud.bigquery.schema import PolicyTagList
685+
686+
return PolicyTagList
687+
688+
def _make_one(self, *args, **kw):
689+
return self._get_target_class()(*args, **kw)
690+
691+
def test_constructor(self):
692+
empty_policy_tags = self._make_one()
693+
self.assertIsNotNone(empty_policy_tags.names)
694+
self.assertEqual(len(empty_policy_tags.names), 0)
695+
policy_tags = self._make_one(["foo", "bar"])
696+
self.assertEqual(policy_tags.names, ("foo", "bar"))
697+
698+
def test_from_api_repr(self):
699+
klass = self._get_target_class()
700+
api_repr = {"names": ["foo"]}
701+
policy_tags = klass.from_api_repr(api_repr)
702+
self.assertEqual(policy_tags.to_api_repr(), api_repr)
703+
704+
# Ensure the None case correctly returns None, rather
705+
# than an empty instance.
706+
policy_tags2 = klass.from_api_repr(None)
707+
self.assertIsNone(policy_tags2)
708+
709+
def test_to_api_repr(self):
710+
taglist = self._make_one(names=["foo", "bar"])
711+
self.assertEqual(
712+
taglist.to_api_repr(), {"names": ["foo", "bar"]},
713+
)
714+
taglist2 = self._make_one(names=("foo", "bar"))
715+
self.assertEqual(
716+
taglist2.to_api_repr(), {"names": ["foo", "bar"]},
717+
)
718+
719+
def test___eq___wrong_type(self):
720+
policy = self._make_one(names=["foo"])
721+
other = object()
722+
self.assertNotEqual(policy, other)
723+
self.assertEqual(policy, mock.ANY)
724+
725+
def test___eq___names_mismatch(self):
726+
policy = self._make_one(names=["foo", "bar"])
727+
other = self._make_one(names=["bar", "baz"])
728+
self.assertNotEqual(policy, other)
729+
730+
def test___hash__set_equality(self):
731+
policy1 = self._make_one(["foo", "bar"])
732+
policy2 = self._make_one(["bar", "baz"])
733+
set_one = {policy1, policy2}
734+
set_two = {policy1, policy2}
735+
self.assertEqual(set_one, set_two)
736+
737+
def test___hash__not_equals(self):
738+
policy1 = self._make_one(["foo", "bar"])
739+
policy2 = self._make_one(["bar", "baz"])
740+
set_one = {policy1}
741+
set_two = {policy2}
742+
self.assertNotEqual(set_one, set_two)

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.