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

Author: plamut

google/cloud/bigquery/schema.py

@@ -15,6 +15,7 @@
 """Schemas for BigQuery tables / queries."""
 
 import collections
+from typing import Optional
 
 from google.cloud.bigquery_v2 import types
 
@@ -105,7 +106,26 @@ def __init__(
  if max_length is not _DEFAULT_VALUE:
  self._properties["maxLength"] = max_length
  self._fields = tuple(fields)
- self._policy_tags = policy_tags
+
+ self._policy_tags = self._determine_policy_tags(field_type, policy_tags)
+
+ @staticmethod
+ def _determine_policy_tags(
+ field_type: str, given_policy_tags: Optional["PolicyTagList"]
+ ) -> Optional["PolicyTagList"]:
+ """Return the given policy tags, or their suitable representation if `None`.
+
+ Args:
+ field_type: The type of the schema field.
+ given_policy_tags: The policy tags to maybe ajdust.
+ """
+ if given_policy_tags is not None:
+ return given_policy_tags
+
+ if field_type is not None and field_type.upper() in _STRUCT_TYPES:
+ return None
+
+ return PolicyTagList()
 
  @staticmethod
  def __get_int(api_repr, name):
@@ -126,18 +146,24 @@ def from_api_repr(cls, api_repr: dict) -> "SchemaField":
  Returns:
  google.cloud.biquery.schema.SchemaField: The ``SchemaField`` object.
  """
+ field_type = api_repr["type"].upper()
+
  # Handle optional properties with default values
  mode = api_repr.get("mode", "NULLABLE")
  description = api_repr.get("description", _DEFAULT_VALUE)
  fields = api_repr.get("fields", ())
 
+ policy_tags = cls._determine_policy_tags(
+ field_type, PolicyTagList.from_api_repr(api_repr.get("policyTags"))
+ )
+
  return cls(
- field_type=api_repr["type"].upper(),
+ field_type=field_type,
  fields=[cls.from_api_repr(f) for f in fields],
  mode=mode.upper(),
  description=description,
  name=api_repr["name"],
- policy_tags=PolicyTagList.from_api_repr(api_repr.get("policyTags")),
+ policy_tags=policy_tags,
  precision=cls.__get_int(api_repr, "precision"),
  scale=cls.__get_int(api_repr, "scale"),
  max_length=cls.__get_int(api_repr, "maxLength"),
@@ -218,9 +244,9 @@ def to_api_repr(self) -> dict:
  # add this to the serialized representation.
  if self.field_type.upper() in _STRUCT_TYPES:
  answer["fields"] = [f.to_api_repr() for f in self.fields]
-
- # If this contains a policy tag definition, include that as well:
- if self.policy_tags is not None:
+ else:
+  # Explicitly include policy tag definition (we must not do it for RECORD
+  # fields, because those are not leaf fields).
  answer["policyTags"] = self.policy_tags.to_api_repr()
 
  # Done; return the serialized dictionary.
@@ -244,14 +270,19 @@ def _key(self):
  field_type = f"{field_type}({self.precision}, {self.scale})"
  else:
  field_type = f"{field_type}({self.precision})"
+
+ policy_tags = (
+ () if self._policy_tags is None else tuple(sorted(self._policy_tags.names))
+ )
+
  return (
  self.name,
  field_type,
  # Mode is always str, if not given it defaults to a str value
  self.mode.upper(), # pytype: disable=attribute-error
  self.description,
  self._fields,
- self._policy_tags,
+ policy_tags,
  )
 
  def to_standard_sql(self) -> types.StandardSqlField:

tests/system/test_client.py

@@ -653,6 +653,56 @@ def test_update_table_schema(self):
  self.assertEqual(found.field_type, expected.field_type)
  self.assertEqual(found.mode, expected.mode)
 
+ def test_unset_table_schema_attributes(self):
+ from google.cloud.bigquery.schema import PolicyTagList
+
+ dataset = self.temp_dataset(_make_dataset_id("unset_policy_tags"))
+ table_id = "test_table"
+ policy_tags = PolicyTagList(
+ names=[
+ "projects/{}/locations/us/taxonomies/1/policyTags/2".format(
+ Config.CLIENT.project
+ ),
+ ]
+ )
+
+ schema = [
+ bigquery.SchemaField("full_name", "STRING", mode="REQUIRED"),
+ bigquery.SchemaField(
+ "secret_int",
+ "INTEGER",
+ mode="REQUIRED",
+ description="This field is numeric",
+ policy_tags=policy_tags,
+ ),
+ ]
+ table_arg = Table(dataset.table(table_id), schema=schema)
+ self.assertFalse(_table_exists(table_arg))
+
+ table = helpers.retry_403(Config.CLIENT.create_table)(table_arg)
+ self.to_delete.insert(0, table)
+
+ self.assertTrue(_table_exists(table))
+ self.assertEqual(policy_tags, table.schema[1].policy_tags)
+
+ # Amend the schema to replace the policy tags
+ new_schema = table.schema[:]
+ old_field = table.schema[1]
+ new_schema[1] = bigquery.SchemaField(
+ name=old_field.name,
+ field_type=old_field.field_type,
+ mode=old_field.mode,
+ description=None,
+ fields=old_field.fields,
+ policy_tags=None,
+ )
+
+ table.schema = new_schema
+ updated_table = Config.CLIENT.update_table(table, ["schema"])
+
+ self.assertFalse(updated_table.schema[1].description) # Empty string or None.
+ self.assertEqual(updated_table.schema[1].policy_tags.names, ())
+
  def test_update_table_clustering_configuration(self):
  dataset = self.temp_dataset(_make_dataset_id("update_table"))
 

tests/unit/job/test_load_config.py

@@ -434,11 +434,13 @@ def test_schema_setter_fields(self):
  "name": "full_name",
  "type": "STRING",
  "mode": "REQUIRED",
+ "policyTags": {"names": []},
  }
  age_repr = {
  "name": "age",
  "type": "INTEGER",
  "mode": "REQUIRED",
+ "policyTags": {"names": []},
  }
  self.assertEqual(
  config._properties["load"]["schema"], {"fields": [full_name_repr, age_repr]}
@@ -451,11 +453,13 @@ def test_schema_setter_valid_mappings_list(self):
  "name": "full_name",
  "type": "STRING",
  "mode": "REQUIRED",
+ "policyTags": {"names": []},
  }
  age_repr = {
  "name": "age",
  "type": "INTEGER",
  "mode": "REQUIRED",
+ "policyTags": {"names": []},
  }
  schema = [full_name_repr, age_repr]
  config.schema = schema

tests/unit/test_client.py

@@ -1019,8 +1019,18 @@ def test_create_table_w_schema_and_query(self):
  {
  "schema": {
  "fields": [
- {"name": "full_name", "type": "STRING", "mode": "REQUIRED"},
- {"name": "age", "type": "INTEGER", "mode": "REQUIRED"},
+ {
+ "name": "full_name",
+ "type": "STRING",
+ "mode": "REQUIRED",
+ "policyTags": {"names": []},
+ },
+ {
+ "name": "age",
+ "type": "INTEGER",
+ "mode": "REQUIRED",
+ "policyTags": {"names": []},
+ },
  ]
  },
  "view": {"query": query},
@@ -1054,8 +1064,18 @@ def test_create_table_w_schema_and_query(self):
  },
  "schema": {
  "fields": [
- {"name": "full_name", "type": "STRING", "mode": "REQUIRED"},
- {"name": "age", "type": "INTEGER", "mode": "REQUIRED"},
+ {
+ "name": "full_name",
+ "type": "STRING",
+ "mode": "REQUIRED",
+ "policyTags": {"names": []},
+ },
+ {
+ "name": "age",
+ "type": "INTEGER",
+ "mode": "REQUIRED",
+ "policyTags": {"names": []},
+ },
  ]
  },
  "view": {"query": query, "useLegacySql": False},
@@ -2000,12 +2020,14 @@ def test_update_table(self):
  "type": "STRING",
  "mode": "REQUIRED",
  "description": None,
+ "policyTags": {"names": []},
  },
  {
  "name": "age",
  "type": "INTEGER",
  "mode": "REQUIRED",
  "description": "New field description",
+ "policyTags": {"names": []},
  },
  ]
  },
@@ -2047,12 +2069,14 @@ def test_update_table(self):
  "type": "STRING",
  "mode": "REQUIRED",
  "description": None,
+ "policyTags": {"names": []},
  },
  {
  "name": "age",
  "type": "INTEGER",
  "mode": "REQUIRED",
  "description": "New field description",
+ "policyTags": {"names": []},
  },
  ]
  },
@@ -2173,14 +2197,21 @@ def test_update_table_w_query(self):
  "type": "STRING",
  "mode": "REQUIRED",
  "description": None,
+ "policyTags": {"names": []},
  },
  {
  "name": "age",
  "type": "INTEGER",
  "mode": "REQUIRED",
  "description": "this is a column",
+ "policyTags": {"names": []},
+ },
+ {
+ "name": "country",
+ "type": "STRING",
+ "mode": "NULLABLE",
+ "policyTags": {"names": []},
  },
- {"name": "country", "type": "STRING", "mode": "NULLABLE"},
  ]
  }
  schema = [
@@ -6516,10 +6547,10 @@ def test_load_table_from_dataframe(self):
  assert field["type"] == table_field.field_type
  assert field["mode"] == table_field.mode
  assert len(field.get("fields", [])) == len(table_field.fields)
+ assert field["policyTags"]["names"] == []
  # Omit unnecessary fields when they come from getting the table
  # (not passed in via job_config)
  assert "description" not in field
- assert "policyTags" not in field
 
  @unittest.skipIf(pandas is None, "Requires `pandas`")
  @unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
@@ -7718,18 +7749,21 @@ def test_schema_to_json_with_file_path(self):
  "description": "quarter",
  "mode": "REQUIRED",
  "name": "qtr",
+ "policyTags": {"names": []},
  "type": "STRING",
  },
  {
  "description": "sales representative",
  "mode": "NULLABLE",
  "name": "rep",
+ "policyTags": {"names": []},
  "type": "STRING",
  },
  {
  "description": "total sales",
  "mode": "NULLABLE",
  "name": "sales",
+ "policyTags": {"names": []},
  "type": "FLOAT",
  },
  ]
@@ -7762,18 +7796,21 @@ def test_schema_to_json_with_file_object(self):
  "description": "quarter",
  "mode": "REQUIRED",
  "name": "qtr",
+ "policyTags": {"names": []},
  "type": "STRING",
  },
  {
  "description": "sales representative",
  "mode": "NULLABLE",
  "name": "rep",
+ "policyTags": {"names": []},
  "type": "STRING",
  },
  {
  "description": "total sales",
  "mode": "NULLABLE",
  "name": "sales",
+ "policyTags": {"names": []},
  "type": "FLOAT",
  },
  ]

tests/unit/test_external_config.py

@@ -78,7 +78,14 @@ def test_to_api_repr_base(self):
  ec.schema = [schema.SchemaField("full_name", "STRING", mode="REQUIRED")]
 
  exp_schema = {
- "fields": [{"name": "full_name", "type": "STRING", "mode": "REQUIRED"}]
+ "fields": [
+ {
+ "name": "full_name",
+ "type": "STRING",
+ "mode": "REQUIRED",
+ "policyTags": {"names": []},
+ }
+ ]
  }
  got_resource = ec.to_api_repr()
  exp_resource = {