From eff365dc17755d0855338e2f273428ffe2056f67 Mon Sep 17 00:00:00 2001 From: Lingqing Gan Date: Wed, 1 Nov 2023 19:49:09 -0400 Subject: [PATCH] feat: support data_governance_type (#1708) * feat: support data_governance_type * remove value validation, add sys test --- google/cloud/bigquery/routine/routine.py | 24 +++++++++++- tests/system/test_client.py | 36 ++++++++++++++++++ tests/unit/routine/test_routine.py | 47 ++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 2 deletions(-) diff --git a/google/cloud/bigquery/routine/routine.py b/google/cloud/bigquery/routine/routine.py index ef33d507e..83cb6362d 100644 --- a/google/cloud/bigquery/routine/routine.py +++ b/google/cloud/bigquery/routine/routine.py @@ -68,6 +68,7 @@ class Routine(object): "description": "description", "determinism_level": "determinismLevel", "remote_function_options": "remoteFunctionOptions", + "data_governance_type": "dataGovernanceType", } def __init__(self, routine_ref, **kwargs) -> None: @@ -300,8 +301,8 @@ def determinism_level(self, value): @property def remote_function_options(self): - """Optional[google.cloud.bigquery.routine.RemoteFunctionOptions]: Configures remote function - options for a routine. + """Optional[google.cloud.bigquery.routine.RemoteFunctionOptions]: + Configures remote function options for a routine. Raises: ValueError: @@ -329,6 +330,25 @@ def remote_function_options(self, value): self._PROPERTY_TO_API_FIELD["remote_function_options"] ] = api_repr + @property + def data_governance_type(self): + """Optional[str]: If set to ``DATA_MASKING``, the function is validated + and made available as a masking function. + + Raises: + ValueError: + If the value is not :data:`string` or :data:`None`. + """ + return self._properties.get(self._PROPERTY_TO_API_FIELD["data_governance_type"]) + + @data_governance_type.setter + def data_governance_type(self, value): + if value is not None and not isinstance(value, str): + raise ValueError( + "invalid data_governance_type, must be a string or `None`." + ) + self._properties[self._PROPERTY_TO_API_FIELD["data_governance_type"]] = value + @classmethod def from_api_repr(cls, resource: dict) -> "Routine": """Factory: construct a routine given its API representation. diff --git a/tests/system/test_client.py b/tests/system/test_client.py index c8ff551ce..7cea8cfa4 100644 --- a/tests/system/test_client.py +++ b/tests/system/test_client.py @@ -13,6 +13,7 @@ # limitations under the License. import base64 +import copy import csv import datetime import decimal @@ -2236,6 +2237,41 @@ def test_create_tvf_routine(self): ] assert result_rows == expected + def test_create_routine_w_data_governance(self): + routine_name = "routine_with_data_governance" + dataset = self.temp_dataset(_make_dataset_id("create_routine")) + + routine = bigquery.Routine( + dataset.routine(routine_name), + type_="SCALAR_FUNCTION", + language="SQL", + body="x", + arguments=[ + bigquery.RoutineArgument( + name="x", + data_type=bigquery.StandardSqlDataType( + type_kind=bigquery.StandardSqlTypeNames.INT64 + ), + ) + ], + data_governance_type="DATA_MASKING", + return_type=bigquery.StandardSqlDataType( + type_kind=bigquery.StandardSqlTypeNames.INT64 + ), + ) + routine_original = copy.deepcopy(routine) + + client = Config.CLIENT + routine_new = client.create_routine(routine) + + assert routine_new.reference == routine_original.reference + assert routine_new.type_ == routine_original.type_ + assert routine_new.language == routine_original.language + assert routine_new.body == routine_original.body + assert routine_new.arguments == routine_original.arguments + assert routine_new.return_type == routine_original.return_type + assert routine_new.data_governance_type == routine_original.data_governance_type + def test_create_table_rows_fetch_nested_schema(self): table_name = "test_table" dataset = self.temp_dataset(_make_dataset_id("create_table_nested_schema")) diff --git a/tests/unit/routine/test_routine.py b/tests/unit/routine/test_routine.py index 87767200c..acd3bc40e 100644 --- a/tests/unit/routine/test_routine.py +++ b/tests/unit/routine/test_routine.py @@ -154,6 +154,7 @@ def test_from_api_repr(target_class): "foo": "bar", }, }, + "dataGovernanceType": "DATA_MASKING", } actual_routine = target_class.from_api_repr(resource) @@ -192,6 +193,7 @@ def test_from_api_repr(target_class): assert actual_routine.remote_function_options.connection == "connection_string" assert actual_routine.remote_function_options.max_batching_rows == 50 assert actual_routine.remote_function_options.user_defined_context == {"foo": "bar"} + assert actual_routine.data_governance_type == "DATA_MASKING" def test_from_api_repr_tvf_function(target_class): @@ -294,6 +296,7 @@ def test_from_api_repr_w_minimal_resource(target_class): assert actual_routine.description is None assert actual_routine.determinism_level is None assert actual_routine.remote_function_options is None + assert actual_routine.data_governance_type is None def test_from_api_repr_w_unknown_fields(target_class): @@ -428,6 +431,20 @@ def test_from_api_repr_w_unknown_fields(target_class): "determinismLevel": bigquery.DeterminismLevel.DETERMINISM_LEVEL_UNSPECIFIED }, ), + ( + { + "arguments": [{"name": "x", "dataType": {"typeKind": "INT64"}}], + "definitionBody": "x * 3", + "language": "SQL", + "returnType": {"typeKind": "INT64"}, + "routineType": "SCALAR_FUNCTION", + "description": "A routine description.", + "determinismLevel": bigquery.DeterminismLevel.DETERMINISM_LEVEL_UNSPECIFIED, + "dataGovernanceType": "DATA_MASKING", + }, + ["data_governance_type"], + {"dataGovernanceType": "DATA_MASKING"}, + ), ( {}, [ @@ -554,6 +571,36 @@ def test_set_remote_function_options_w_none(object_under_test): assert object_under_test._properties["remoteFunctionOptions"] is None +def test_set_data_governance_type_w_none(object_under_test): + object_under_test.data_governance_type = None + assert object_under_test.data_governance_type is None + assert object_under_test._properties["dataGovernanceType"] is None + + +def test_set_data_governance_type_valid(object_under_test): + object_under_test.data_governance_type = "DATA_MASKING" + assert object_under_test.data_governance_type == "DATA_MASKING" + assert object_under_test._properties["dataGovernanceType"] == "DATA_MASKING" + + +def test_set_data_governance_type_wrong_type(object_under_test): + with pytest.raises(ValueError) as exp: + object_under_test.data_governance_type = 1 + assert "invalid data_governance_type" in str(exp) + assert object_under_test.data_governance_type is None + assert object_under_test._properties.get("dataGovernanceType") is None + + +def test_set_data_governance_type_wrong_str(object_under_test): + """Client does not verify the content of data_governance_type string to be + compatible with future upgrades. If the value is not supported, BigQuery + itself will report an error. + """ + object_under_test.data_governance_type = "RANDOM_STRING" + assert object_under_test.data_governance_type == "RANDOM_STRING" + assert object_under_test._properties["dataGovernanceType"] == "RANDOM_STRING" + + def test_repr(target_class): model = target_class("my-proj.my_dset.my_routine") actual_routine = repr(model)