feat(bigquery): add column data policies to SchemaField

packages/google-cloud-bigquery/docs/usage/tables.rst

@@ -183,6 +183,39 @@ Load a JSON file from Cloud Storage and auto-detect schema:
    :start-after: [START bigquery_load_table_gcs_json_autodetect]
    :end-before: [END bigquery_load_table_gcs_json_autodetect]
 
+Column Data Policies
+^^^^^^^^^^^^^^^^^^^^
+
+Use :class:`~google.cloud.bigquery.schema.SchemaField` to represent data policies
+attached directly to a column. Pass the policy resource names as
+``data_policies``:
+
+.. code-block:: python
+
+    from google.cloud import bigquery
+
+    field = bigquery.SchemaField(
+        "ssn",
+        "STRING",
+        data_policies=["projects/my-project/locations/us/dataPolicies/mask_ssn"],
+    )
+
+The ``data_policies`` property returns a tuple of resource names, or ``None``
+when the property is unset. The client serializes these names as ``dataPolicies``
+objects in the table schema. Omitting the argument leaves that property out of
+the API representation; an empty iterable produces an empty array.
+
+This property represents ``dataPolicies``, not the separate ``dataPolicyList``
+REST field. When ``dataPolicyList`` is present, BigQuery uses it as the source of
+truth for data policy information.
+
+BigQuery does not allow policy tags and data policies on the same column. To
+remove the last data policy from a column, use ``ALTER COLUMN SET OPTIONS`` with
+``data_policies = []``. Serializing an empty array does not change the service's
+policy-removal behavior. See `Mask data with data policies directly on a column
+<https://cloud.google.com/bigquery/docs/column-data-masking>`_
+for permissions and service limitations.
+
 Updating a Table
 ^^^^^^^^^^^^^^^^
 

packages/google-cloud-bigquery/google/cloud/bigquery/schema.py

@@ -143,6 +143,7 @@ class _SchemaFieldKey(typing.NamedTuple):
     timestamp_precision: Optional[int]
     rounding_mode: Optional[str]
     foreign_type_definition: Optional[str]
+    data_policies: Optional[tuple[str, ...]]
 
 
 class SchemaField(object):
@@ -239,6 +240,15 @@ class SchemaField(object):
             Defaults to `enums.TimestampPrecision.MICROSECOND` (`None`) for
             microsecond precision. Use `enums.TimestampPrecision.PICOSECOND`
             (`12`) for picosecond precision.
+
+        data_policies: Optional[Iterable[str]]
+            Data policy resource names attached directly to the field, in the
+            form ``projects/{project}/locations/{location}/dataPolicies/{policy}``.
+            Serialized as the REST ``dataPolicies`` field. If omitted, the
+            property is not included in the API representation.
+            ``None`` is serialized as ``null`` and an empty iterable as ``[]``.
+            A single ``str`` or ``bytes`` value raises :class:`TypeError`;
+            use an iterable containing the name instead.
     """
 
     def __init__(
@@ -257,6 +267,7 @@ def __init__(
         rounding_mode: Union[enums.RoundingMode, str, None] = None,
         foreign_type_definition: Optional[str] = None,
         timestamp_precision: Optional[enums.TimestampPrecision] = None,
+        data_policies: Union[Iterable[str], None, _DefaultSentinel] = _DEFAULT_VALUE,
     ):
         self._properties: Dict[str, Any] = {
             "name": name,
@@ -296,6 +307,16 @@ def __init__(
             self._properties["roundingMode"] = rounding_mode
         if foreign_type_definition is not None:
             self._properties["foreignTypeDefinition"] = foreign_type_definition
+        if data_policies is not _DEFAULT_VALUE:
+            if isinstance(data_policies, (str, bytes)):
+                raise TypeError(
+                    "data_policies must be an iterable of strings, not str or bytes"
+                )
+            self._properties["dataPolicies"] = (
+                [{"name": name} for name in data_policies]
+                if data_policies is not None
+                else None
+            )
 
         if fields:  # Don't set the property if it's not set.
             self._properties["fields"] = [field.to_api_repr() for field in fields]
@@ -432,6 +453,23 @@ def policy_tags(self):
         resource = self._properties.get("policyTags")
         return PolicyTagList.from_api_repr(resource) if resource is not None else None
 
+    @property
+    def data_policies(self) -> Optional[tuple[str, ...]]:
+        """Optional[Tuple[str, ...]]: Data policy resource names for this field.
+
+        Returns the names from the REST ``dataPolicies`` field, or ``None`` if
+        unset. An explicitly empty list is returned as an empty tuple.
+
+        See:
+            https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#TableFieldSchema
+        """
+        resource = self._properties.get("dataPolicies")
+        return (
+            tuple(policy["name"] for policy in resource)
+            if resource is not None
+            else None
+        )
+
     @property
     def timestamp_precision(self) -> enums.TimestampPrecision:
         """Precision (maximum number of total digits in base 10) for seconds of
@@ -478,6 +516,7 @@ def _key(self):
         policy_tags = (
             None if self.policy_tags is None else tuple(sorted(self.policy_tags.names))
         )
+        data_policies = self.data_policies
 
         timestamp_precision = self._properties.get("timestampPrecision")
 
@@ -499,6 +538,9 @@ def _key(self):
             timestamp_precision=timestamp_precision,
             rounding_mode=rounding_mode,
             foreign_type_definition=self.foreign_type_definition,
+            data_policies=(
+                None if data_policies is None else tuple(sorted(data_policies))
+            ),
         )
 
     def to_standard_sql(self) -> standard_sql.StandardSqlField:
@@ -553,8 +595,19 @@ def __repr__(self):
         policy_tags_inst = (
             None if key.policy_tags is None else PolicyTagList(key.policy_tags)
         )
-        adjusted_key = tuple(key._replace(policy_tags=policy_tags_inst))
-        return f"{self.__class__.__name__}{adjusted_key}"
+        # Keep data_policies keyword-based without changing the existing repr.
+        key_to_show = key._replace(policy_tags=policy_tags_inst)
+        adjusted_key = tuple(
+            getattr(key_to_show, field)
+            for field in key_to_show._fields
+            if field != "data_policies"
+        )
+        representation = f"{self.__class__.__name__}{adjusted_key}"
+        if key.data_policies is not None:
+            representation = (
+                f"{representation[:-1]}, data_policies={key.data_policies!r})"
+            )
+        return representation
 
 
 def _parse_schema_resource(info):

packages/google-cloud-bigquery/tests/unit/test_schema.py

@@ -48,6 +48,7 @@ def test_constructor_defaults(self):
         self.assertIsNone(field.description)
         self.assertEqual(field.fields, ())
         self.assertIsNone(field.policy_tags)
+        self.assertIsNone(field.data_policies)
         self.assertIsNone(field.default_value_expression)
         self.assertEqual(field.rounding_mode, None)
         self.assertEqual(field.foreign_type_definition, None)
@@ -99,6 +100,7 @@ def test_constructor_explicit_none(self):
         field = self._make_one("test", "STRING", description=None, policy_tags=None)
         self.assertIsNone(field.description)
         self.assertIsNone(field.policy_tags)
+        self.assertIsNone(field.data_policies)
 
     def test_constructor_subfields(self):
         sub_field1 = self._make_one("area_code", "STRING")
@@ -181,6 +183,7 @@ def test_to_api_repr_omits_unset_properties(self):
         resource = field.to_api_repr()
         self.assertNotIn("description", resource)
         self.assertNotIn("policyTags", resource)
+        self.assertNotIn("dataPolicies", resource)
 
     def test_to_api_repr_with_subfield(self):
         for record_type in ("RECORD", "STRUCT"):
@@ -288,9 +291,11 @@ def test_from_api_repr_defaults(self):
         # _properties.
         self.assertIsNone(field.description)
         self.assertIsNone(field.policy_tags)
+        self.assertIsNone(field.data_policies)
         self.assertIsNone(field.range_element_type)
         self.assertNotIn("description", field._properties)
         self.assertNotIn("policyTags", field._properties)
+        self.assertNotIn("dataPolicies", field._properties)
         self.assertNotIn("rangeElementType", field._properties)
 
     def test_from_api_repr_timestamp_precision_str(self):
@@ -763,6 +768,198 @@ def test___repr__evaluable_with_policy_tags(self):
         assert field == evaled_field
 
 
+class TestSchemaFieldDataPolicies:
+    POLICY_ONE = "projects/project/locations/us/dataPolicies/one"
+    POLICY_TWO = "projects/project/locations/us/dataPolicies/two"
+
+    @pytest.mark.parametrize(
+        "names, expected",
+        [
+            (None, None),
+            ([], []),
+            ([POLICY_ONE], [{"name": POLICY_ONE}]),
+            (
+                [POLICY_ONE, POLICY_TWO],
+                [{"name": POLICY_ONE}, {"name": POLICY_TWO}],
+            ),
+        ],
+    )
+    def test_constructor_and_round_trip(self, names, expected):
+        field = schema.SchemaField("ssn", "STRING", data_policies=names)
+        resource = {
+            "name": "ssn",
+            "type": "STRING",
+            "mode": "NULLABLE",
+            "dataPolicies": expected,
+        }
+
+        assert field.to_api_repr() == resource
+        assert field.data_policies == (None if names is None else tuple(names))
+        restored = schema.SchemaField.from_api_repr(copy.deepcopy(resource))
+        assert restored == field
+        assert hash(restored) == hash(field)
+        assert restored.data_policies == field.data_policies
+        assert restored.to_api_repr() == resource
+
+    @pytest.mark.parametrize(
+        "properties, expected",
+        [
+            ({}, None),
+            ({"dataPolicies": None}, None),
+            ({"dataPolicies": []}, ()),
+            ({"dataPolicies": [{"name": POLICY_ONE}]}, (POLICY_ONE,)),
+            (
+                {"dataPolicies": [{"name": POLICY_ONE}, {"name": POLICY_TWO}]},
+                (POLICY_ONE, POLICY_TWO),
+            ),
+        ],
+    )
+    def test_from_api_repr(self, properties, expected):
+        resource = {"name": "ssn", "type": "STRING", "mode": "NULLABLE"}
+        resource.update(properties)
+        original = copy.deepcopy(resource)
+        field = schema.SchemaField.from_api_repr(resource)
+
+        assert field.data_policies == expected
+        assert field.to_api_repr() == original
+        assert field.to_api_repr() is resource
+
+    @pytest.mark.parametrize("container", [list, tuple, iter])
+    def test_constructor_accepts_iterables(self, container):
+        names = container([self.POLICY_ONE, self.POLICY_TWO])
+        field = schema.SchemaField("ssn", "STRING", data_policies=names)
+
+        assert field.data_policies == (self.POLICY_ONE, self.POLICY_TWO)
+        assert field.to_api_repr()["dataPolicies"] == [
+            {"name": self.POLICY_ONE},
+            {"name": self.POLICY_TWO},
+        ]
+
+    @pytest.mark.parametrize("names", [POLICY_ONE, POLICY_ONE.encode(), "", b""])
+    def test_constructor_rejects_strings_and_bytes(self, names):
+        with pytest.raises(TypeError, match="data_policies must be an iterable"):
+            schema.SchemaField("ssn", "STRING", data_policies=names)
+
+    def test_constructor_copies_names(self):
+        names = [self.POLICY_ONE]
+        field = schema.SchemaField("ssn", "STRING", data_policies=names)
+        names.append(self.POLICY_TWO)
+
+        assert field.data_policies == (self.POLICY_ONE,)
+
+    def test_equality_ignores_order(self):
+        field = schema.SchemaField(
+            "ssn", "STRING", data_policies=[self.POLICY_ONE, self.POLICY_TWO]
+        )
+        other = schema.SchemaField(
+            "ssn", "STRING", data_policies=[self.POLICY_TWO, self.POLICY_ONE]
+        )
+
+        assert field == other
+        assert not field != other
+        assert hash(field) == hash(other)
+        assert len({field, other}) == 1
+        assert field._key().data_policies == (self.POLICY_ONE, self.POLICY_TWO)
+        # Comparing fields must not reorder the API representation.
+        assert other.data_policies == (self.POLICY_TWO, self.POLICY_ONE)
+
+    @pytest.mark.parametrize("other_names", [None, [], [POLICY_TWO]])
+    def test_different_policies_are_not_equal(self, other_names):
+        field = schema.SchemaField("ssn", "STRING", data_policies=[self.POLICY_ONE])
+        other = schema.SchemaField("ssn", "STRING", data_policies=other_names)
+
+        assert field != other
+        assert len({field, other}) == 2
+        assert field._key().data_policies != other._key().data_policies
+
+    def test_absent_and_none_are_equal(self):
+        field = schema.SchemaField("ssn", "STRING")
+        other = schema.SchemaField("ssn", "STRING", data_policies=None)
+
+        assert field == other
+        assert hash(field) == hash(other)
+
+    def test_absent_and_empty_are_not_equal(self):
+        field = schema.SchemaField("ssn", "STRING")
+        other = schema.SchemaField("ssn", "STRING", data_policies=[])
+
+        assert field != other
+
+    def test_nested_field_round_trip(self):
+        child = schema.SchemaField("ssn", "STRING", data_policies=[self.POLICY_ONE])
+        field = schema.SchemaField("person", "RECORD", fields=[child])
+        resource = {
+            "name": "person",
+            "type": "RECORD",
+            "mode": "NULLABLE",
+            "fields": [
+                {
+                    "name": "ssn",
+                    "type": "STRING",
+                    "mode": "NULLABLE",
+                    "dataPolicies": [{"name": self.POLICY_ONE}],
+                }
+            ],
+        }
+
+        assert field.to_api_repr() == resource
+        restored = schema.SchemaField.from_api_repr(copy.deepcopy(resource))
+        assert restored.fields[0].data_policies == (self.POLICY_ONE,)
+        assert restored == field
+        assert hash(restored) == hash(field)
+
+        other_child = schema.SchemaField(
+            "ssn", "STRING", data_policies=[self.POLICY_TWO]
+        )
+        other = schema.SchemaField("person", "RECORD", fields=[other_child])
+        assert field != other
+
+    def test_policy_tags_on_another_field(self):
+        tagged = schema.SchemaField(
+            "name", "STRING", policy_tags=PolicyTagList(["policy-tag"])
+        )
+        masked = schema.SchemaField("ssn", "STRING", data_policies=[self.POLICY_ONE])
+        fields = schema._to_schema_fields(
+            schema._build_schema_resource([tagged, masked])
+        )
+
+        assert fields == [tagged, masked]
+        assert fields[0].policy_tags.names == ("policy-tag",)
+        assert fields[0].data_policies is None
+        assert "dataPolicies" not in fields[0].to_api_repr()
+        assert fields[1].policy_tags is None
+        assert "policyTags" not in fields[1].to_api_repr()
+        assert fields[1].data_policies == (self.POLICY_ONE,)
+
+    @pytest.mark.parametrize("names", [[], [POLICY_ONE], [POLICY_TWO, POLICY_ONE]])
+    def test_repr(self, names):
+        field = schema.SchemaField("ssn", "STRING", data_policies=names)
+
+        restored = eval(repr(field), {"SchemaField": schema.SchemaField})
+        assert restored == field
+        assert restored.data_policies == tuple(sorted(names))
+
+    def test_repr_with_nested_fields(self):
+        field = schema.SchemaField(
+            "person",
+            "RECORD",
+            fields=[
+                schema.SchemaField(
+                    "name", "STRING", policy_tags=PolicyTagList(["policy-tag"])
+                ),
+                schema.SchemaField("ssn", "STRING", data_policies=[self.POLICY_ONE]),
+            ],
+        )
+
+        restored = eval(
+            repr(field),
+            {"SchemaField": schema.SchemaField, "PolicyTagList": PolicyTagList},
+        )
+        assert restored == field
+        assert restored.fields[0].policy_tags == field.fields[0].policy_tags
+        assert restored.fields[1].data_policies == (self.POLICY_ONE,)
+
+
 class TestFieldElementType(unittest.TestCase):
     @staticmethod
     def _get_target_class():