diff --git a/google/cloud/bigquery/__init__.py b/google/cloud/bigquery/__init__.py index ebd5b3109..40e3a1578 100644 --- a/google/cloud/bigquery/__init__.py +++ b/google/cloud/bigquery/__init__.py @@ -93,6 +93,7 @@ from google.cloud.bigquery.routine import RoutineArgument from google.cloud.bigquery.routine import RoutineReference from google.cloud.bigquery.routine import RoutineType +from google.cloud.bigquery.routine import RemoteFunctionOptions from google.cloud.bigquery.schema import PolicyTagList from google.cloud.bigquery.schema import SchemaField from google.cloud.bigquery.standard_sql import StandardSqlDataType @@ -154,6 +155,7 @@ "Routine", "RoutineArgument", "RoutineReference", + "RemoteFunctionOptions", # Shared helpers "SchemaField", "PolicyTagList", diff --git a/google/cloud/bigquery/routine/__init__.py b/google/cloud/bigquery/routine/__init__.py index 7353073c8..e576b0d49 100644 --- a/google/cloud/bigquery/routine/__init__.py +++ b/google/cloud/bigquery/routine/__init__.py @@ -20,6 +20,7 @@ from google.cloud.bigquery.routine.routine import RoutineArgument from google.cloud.bigquery.routine.routine import RoutineReference from google.cloud.bigquery.routine.routine import RoutineType +from google.cloud.bigquery.routine.routine import RemoteFunctionOptions __all__ = ( @@ -28,4 +29,5 @@ "RoutineArgument", "RoutineReference", "RoutineType", + "RemoteFunctionOptions", ) diff --git a/google/cloud/bigquery/routine/routine.py b/google/cloud/bigquery/routine/routine.py index 3c0919003..36ed03728 100644 --- a/google/cloud/bigquery/routine/routine.py +++ b/google/cloud/bigquery/routine/routine.py @@ -67,6 +67,7 @@ class Routine(object): "type_": "routineType", "description": "description", "determinism_level": "determinismLevel", + "remote_function_options": "remoteFunctionOptions", } def __init__(self, routine_ref, **kwargs) -> None: @@ -297,6 +298,37 @@ def determinism_level(self): def determinism_level(self, value): self._properties[self._PROPERTY_TO_API_FIELD["determinism_level"]] = value + @property + def remote_function_options(self): + """Optional[google.cloud.bigquery.routine.RemoteFunctionOptions]: Configures remote function + options for a routine. + + Raises: + ValueError: + If the value is not + :class:`~google.cloud.bigquery.routine.RemoteFunctionOptions` or + :data:`None`. + """ + prop = self._properties.get( + self._PROPERTY_TO_API_FIELD["remote_function_options"] + ) + if prop is not None: + return RemoteFunctionOptions.from_api_repr(prop) + + @remote_function_options.setter + def remote_function_options(self, value): + api_repr = value + if isinstance(value, RemoteFunctionOptions): + api_repr = value.to_api_repr() + elif value is not None: + raise ValueError( + "value must be google.cloud.bigquery.routine.RemoteFunctionOptions " + "or None" + ) + self._properties[ + self._PROPERTY_TO_API_FIELD["remote_function_options"] + ] = api_repr + @classmethod def from_api_repr(cls, resource: dict) -> "Routine": """Factory: construct a routine given its API representation. @@ -563,3 +595,124 @@ def __str__(self): This is a fully-qualified ID, including the project ID and dataset ID. """ return "{}.{}.{}".format(self.project, self.dataset_id, self.routine_id) + + +class RemoteFunctionOptions(object): + """Configuration options for controlling remote BigQuery functions.""" + + _PROPERTY_TO_API_FIELD = { + "endpoint": "endpoint", + "connection": "connection", + "max_batching_rows": "maxBatchingRows", + "user_defined_context": "userDefinedContext", + } + + def __init__( + self, + endpoint=None, + connection=None, + max_batching_rows=None, + user_defined_context=None, + _properties=None, + ) -> None: + if _properties is None: + _properties = {} + self._properties = _properties + + if endpoint is not None: + self.endpoint = endpoint + if connection is not None: + self.connection = connection + if max_batching_rows is not None: + self.max_batching_rows = max_batching_rows + if user_defined_context is not None: + self.user_defined_context = user_defined_context + + @property + def connection(self): + """string: Fully qualified name of the user-provided connection object which holds the authentication information to send requests to the remote service. + + Format is "projects/{projectId}/locations/{locationId}/connections/{connectionId}" + """ + return _helpers._str_or_none(self._properties.get("connection")) + + @connection.setter + def connection(self, value): + self._properties["connection"] = _helpers._str_or_none(value) + + @property + def endpoint(self): + """string: Endpoint of the user-provided remote service + + Example: "https://us-east1-my_gcf_project.cloudfunctions.net/remote_add" + """ + return _helpers._str_or_none(self._properties.get("endpoint")) + + @endpoint.setter + def endpoint(self, value): + self._properties["endpoint"] = _helpers._str_or_none(value) + + @property + def max_batching_rows(self): + """int64: Max number of rows in each batch sent to the remote service. + + If absent or if 0, BigQuery dynamically decides the number of rows in a batch. + """ + return _helpers._int_or_none(self._properties.get("maxBatchingRows")) + + @max_batching_rows.setter + def max_batching_rows(self, value): + self._properties["maxBatchingRows"] = _helpers._str_or_none(value) + + @property + def user_defined_context(self): + """Dict[str, str]: User-defined context as a set of key/value pairs, + which will be sent as function invocation context together with + batched arguments in the requests to the remote service. The total + number of bytes of keys and values must be less than 8KB. + """ + return self._properties.get("userDefinedContext") + + @user_defined_context.setter + def user_defined_context(self, value): + if not isinstance(value, dict): + raise ValueError("value must be dictionary") + self._properties["userDefinedContext"] = value + + @classmethod + def from_api_repr(cls, resource: dict) -> "RemoteFunctionOptions": + """Factory: construct remote function options given its API representation. + + Args: + resource (Dict[str, object]): Resource, as returned from the API. + + Returns: + google.cloud.bigquery.routine.RemoteFunctionOptions: + Python object, as parsed from ``resource``. + """ + ref = cls() + ref._properties = resource + return ref + + def to_api_repr(self) -> dict: + """Construct the API resource representation of this RemoteFunctionOptions. + + Returns: + Dict[str, object]: Remote function options represented as an API resource. + """ + return self._properties + + def __eq__(self, other): + if not isinstance(other, RemoteFunctionOptions): + return NotImplemented + return self._properties == other._properties + + def __ne__(self, other): + return not self == other + + def __repr__(self): + all_properties = [ + "{}={}".format(property_name, repr(getattr(self, property_name))) + for property_name in sorted(self._PROPERTY_TO_API_FIELD) + ] + return "RemoteFunctionOptions({})".format(", ".join(all_properties)) diff --git a/tests/unit/routine/test_remote_function_options.py b/tests/unit/routine/test_remote_function_options.py new file mode 100644 index 000000000..b476dca1e --- /dev/null +++ b/tests/unit/routine/test_remote_function_options.py @@ -0,0 +1,128 @@ +# -*- coding: utf-8 -*- +# +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +ENDPOINT = "https://some.endpoint" +CONNECTION = "connection_string" +MAX_BATCHING_ROWS = 50 +USER_DEFINED_CONTEXT = { + "foo": "bar", +} + + +@pytest.fixture +def target_class(): + from google.cloud.bigquery.routine import RemoteFunctionOptions + + return RemoteFunctionOptions + + +def test_ctor(target_class): + + options = target_class( + endpoint=ENDPOINT, + connection=CONNECTION, + max_batching_rows=MAX_BATCHING_ROWS, + user_defined_context=USER_DEFINED_CONTEXT, + ) + assert options.endpoint == ENDPOINT + assert options.connection == CONNECTION + assert options.max_batching_rows == MAX_BATCHING_ROWS + assert options.user_defined_context == USER_DEFINED_CONTEXT + + +def test_empty_ctor(target_class): + options = target_class() + assert options._properties == {} + options = target_class(_properties=None) + assert options._properties == {} + options = target_class(_properties={}) + assert options._properties == {} + + +def test_ctor_bad_context(target_class): + with pytest.raises(ValueError, match="value must be dictionary"): + target_class(user_defined_context=[1, 2, 3, 4]) + + +def test_from_api_repr(target_class): + resource = { + "endpoint": ENDPOINT, + "connection": CONNECTION, + "maxBatchingRows": MAX_BATCHING_ROWS, + "userDefinedContext": USER_DEFINED_CONTEXT, + "someRandomField": "someValue", + } + options = target_class.from_api_repr(resource) + assert options.endpoint == ENDPOINT + assert options.connection == CONNECTION + assert options.max_batching_rows == MAX_BATCHING_ROWS + assert options.user_defined_context == USER_DEFINED_CONTEXT + assert options._properties["someRandomField"] == "someValue" + + +def test_from_api_repr_w_minimal_resource(target_class): + resource = {} + options = target_class.from_api_repr(resource) + assert options.endpoint is None + assert options.connection is None + assert options.max_batching_rows is None + assert options.user_defined_context is None + + +def test_from_api_repr_w_unknown_fields(target_class): + resource = {"thisFieldIsNotInTheProto": "just ignore me"} + options = target_class.from_api_repr(resource) + assert options._properties is resource + + +def test_eq(target_class): + options = target_class( + endpoint=ENDPOINT, + connection=CONNECTION, + max_batching_rows=MAX_BATCHING_ROWS, + user_defined_context=USER_DEFINED_CONTEXT, + ) + other_options = target_class( + endpoint=ENDPOINT, + connection=CONNECTION, + max_batching_rows=MAX_BATCHING_ROWS, + user_defined_context=USER_DEFINED_CONTEXT, + ) + assert options == other_options + assert not (options != other_options) + + empty_options = target_class() + assert not (options == empty_options) + assert options != empty_options + + notanarg = object() + assert not (options == notanarg) + assert options != notanarg + + +def test_repr(target_class): + options = target_class( + endpoint=ENDPOINT, + connection=CONNECTION, + max_batching_rows=MAX_BATCHING_ROWS, + user_defined_context=USER_DEFINED_CONTEXT, + ) + actual_repr = repr(options) + assert actual_repr == ( + "RemoteFunctionOptions(connection='connection_string', endpoint='https://some.endpoint', max_batching_rows=50, user_defined_context={'foo': 'bar'})" + ) diff --git a/tests/unit/routine/test_routine.py b/tests/unit/routine/test_routine.py index 80a3def73..87767200c 100644 --- a/tests/unit/routine/test_routine.py +++ b/tests/unit/routine/test_routine.py @@ -75,6 +75,13 @@ def test_ctor_w_properties(target_class): description = "A routine description." determinism_level = bigquery.DeterminismLevel.NOT_DETERMINISTIC + options = bigquery.RemoteFunctionOptions( + endpoint="https://some.endpoint", + connection="connection_string", + max_batching_rows=99, + user_defined_context={"foo": "bar"}, + ) + actual_routine = target_class( routine_id, arguments=arguments, @@ -84,6 +91,7 @@ def test_ctor_w_properties(target_class): type_=type_, description=description, determinism_level=determinism_level, + remote_function_options=options, ) ref = RoutineReference.from_string(routine_id) @@ -97,6 +105,18 @@ def test_ctor_w_properties(target_class): assert ( actual_routine.determinism_level == bigquery.DeterminismLevel.NOT_DETERMINISTIC ) + assert actual_routine.remote_function_options == options + + +def test_ctor_invalid_remote_function_options(target_class): + with pytest.raises( + ValueError, + match=".*must be google.cloud.bigquery.routine.RemoteFunctionOptions.*", + ): + target_class( + "my-proj.my_dset.my_routine", + remote_function_options=object(), + ) def test_from_api_repr(target_class): @@ -126,6 +146,14 @@ def test_from_api_repr(target_class): "someNewField": "someValue", "description": "A routine description.", "determinismLevel": bigquery.DeterminismLevel.DETERMINISTIC, + "remoteFunctionOptions": { + "endpoint": "https://some.endpoint", + "connection": "connection_string", + "maxBatchingRows": 50, + "userDefinedContext": { + "foo": "bar", + }, + }, } actual_routine = target_class.from_api_repr(resource) @@ -160,6 +188,10 @@ def test_from_api_repr(target_class): assert actual_routine._properties["someNewField"] == "someValue" assert actual_routine.description == "A routine description." assert actual_routine.determinism_level == "DETERMINISTIC" + assert actual_routine.remote_function_options.endpoint == "https://some.endpoint" + 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"} def test_from_api_repr_tvf_function(target_class): @@ -261,6 +293,7 @@ def test_from_api_repr_w_minimal_resource(target_class): assert actual_routine.type_ is None assert actual_routine.description is None assert actual_routine.determinism_level is None + assert actual_routine.remote_function_options is None def test_from_api_repr_w_unknown_fields(target_class): @@ -421,6 +454,24 @@ def test_from_api_repr_w_unknown_fields(target_class): ["someNewField"], {"someNewField": "someValue"}, ), + ( + { + "routineType": "SCALAR_FUNCTION", + "remoteFunctionOptions": { + "endpoint": "https://some_endpoint", + "connection": "connection_string", + "max_batching_rows": 101, + }, + }, + ["remote_function_options"], + { + "remoteFunctionOptions": { + "endpoint": "https://some_endpoint", + "connection": "connection_string", + "max_batching_rows": 101, + }, + }, + ), ], ) def test_build_resource(object_under_test, resource, filter_fields, expected): @@ -497,6 +548,12 @@ def test_set_description_w_none(object_under_test): assert object_under_test._properties["description"] is None +def test_set_remote_function_options_w_none(object_under_test): + object_under_test.remote_function_options = None + assert object_under_test.remote_function_options is None + assert object_under_test._properties["remoteFunctionOptions"] is None + + def test_repr(target_class): model = target_class("my-proj.my_dset.my_routine") actual_routine = repr(model)