Skip to content

Commit eed73a1

Browse files
feat: Add ExternalRuntimeOptions to BigQuery routine
This change introduces the `ExternalRuntimeOptions` class to the `google.cloud.bigquery.routine` module, allowing users to configure runtime options for external routines. Key changes: - Created the `ExternalRuntimeOptions` class with setters and getters for `container_memory`, `container_cpu`, `runtime_connection`, `max_batching_rows`, and `runtime_version`. - Updated the `Routine` class to include an `external_runtime_options` property that accepts an `ExternalRuntimeOptions` object. - Added comprehensive unit tests for the new class and its integration with the `Routine` class, including tests for both valid and invalid input values.
1 parent 8bbd3d0 commit eed73a1

File tree

5 files changed

+377
-0
lines changed

5 files changed

+377
-0
lines changed

google/cloud/bigquery/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@
9898
from google.cloud.bigquery.routine import RoutineReference
9999
from google.cloud.bigquery.routine import RoutineType
100100
from google.cloud.bigquery.routine import RemoteFunctionOptions
101+
from google.cloud.bigquery.routine import ExternalRuntimeOptions
101102
from google.cloud.bigquery.schema import PolicyTagList
102103
from google.cloud.bigquery.schema import SchemaField
103104
from google.cloud.bigquery.schema import FieldElementType
@@ -181,6 +182,7 @@
181182
"RoutineArgument",
182183
"RoutineReference",
183184
"RemoteFunctionOptions",
185+
"ExternalRuntimeOptions",
184186
# Shared helpers
185187
"SchemaField",
186188
"FieldElementType",

google/cloud/bigquery/routine/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from google.cloud.bigquery.routine.routine import RoutineReference
2222
from google.cloud.bigquery.routine.routine import RoutineType
2323
from google.cloud.bigquery.routine.routine import RemoteFunctionOptions
24+
from google.cloud.bigquery.routine.routine import ExternalRuntimeOptions
2425

2526

2627
__all__ = (
@@ -30,4 +31,5 @@
3031
"RoutineReference",
3132
"RoutineType",
3233
"RemoteFunctionOptions",
34+
"ExternalRuntimeOptions",
3335
)

google/cloud/bigquery/routine/routine.py

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ class Routine(object):
6969
"determinism_level": "determinismLevel",
7070
"remote_function_options": "remoteFunctionOptions",
7171
"data_governance_type": "dataGovernanceType",
72+
"external_runtime_options": "externalRuntimeOptions",
7273
}
7374

7475
def __init__(self, routine_ref, **kwargs) -> None:
@@ -349,6 +350,37 @@ def data_governance_type(self, value):
349350
)
350351
self._properties[self._PROPERTY_TO_API_FIELD["data_governance_type"]] = value
351352

353+
@property
354+
def external_runtime_options(self):
355+
"""Optional[google.cloud.bigquery.routine.ExternalRuntimeOptions]:
356+
Configures the external runtime options for a routine.
357+
358+
Raises:
359+
ValueError:
360+
If the value is not
361+
:class:`~google.cloud.bigquery.routine.ExternalRuntimeOptions` or
362+
:data:`None`.
363+
"""
364+
prop = self._properties.get(
365+
self._PROPERTY_TO_API_FIELD["external_runtime_options"]
366+
)
367+
if prop is not None:
368+
return ExternalRuntimeOptions.from_api_repr(prop)
369+
370+
@external_runtime_options.setter
371+
def external_runtime_options(self, value):
372+
api_repr = value
373+
if isinstance(value, ExternalRuntimeOptions):
374+
api_repr = value.to_api_repr()
375+
elif value is not None:
376+
raise ValueError(
377+
"value must be google.cloud.bigquery.routine.ExternalRuntimeOptions "
378+
"or None"
379+
)
380+
self._properties[
381+
self._PROPERTY_TO_API_FIELD["external_runtime_options"]
382+
] = api_repr
383+
352384
@classmethod
353385
def from_api_repr(cls, resource: dict) -> "Routine":
354386
"""Factory: construct a routine given its API representation.
@@ -736,3 +768,151 @@ def __repr__(self):
736768
for property_name in sorted(self._PROPERTY_TO_API_FIELD)
737769
]
738770
return "RemoteFunctionOptions({})".format(", ".join(all_properties))
771+
772+
773+
class ExternalRuntimeOptions(object):
774+
"""Options for the runtime of the external system.
775+
Args:
776+
container_memory (str):
777+
Optional. Amount of memory provisioned for a Python UDF container
778+
instance. Format: {number}{unit} where unit is one of "M", "G", "Mi"
779+
and "Gi" (e.g. 1G, 512Mi). If not specified, the default value is
780+
512Mi. For more information, see `Configure container limits for
781+
Python UDFs <https://cloud.google.com/bigquery/docs/user-defined-functions-python#configure-container-limits>`_
782+
container_cpu (int):
783+
Optional. Amount of CPU provisioned for a Python UDF container
784+
instance. For more information, see `Configure container limits
785+
for Python UDFs <https://cloud.google.com/bigquery/docs/user-defined-functions-python#configure-container-limits>`_
786+
runtime_connection (str):
787+
Optional. Fully qualified name of the connection whose service account
788+
will be used to execute the code in the container. Format:
789+
"projects/{projectId}/locations/{locationId}/connections/{connectionId}"
790+
max_batching_rows (int):
791+
Optional. Maximum number of rows in each batch sent to the external
792+
runtime. If absent or if 0, BigQuery dynamically decides the number of
793+
rows in a batch.
794+
runtime_version (str):
795+
Optional. Language runtime version. Example: python-3.11.
796+
"""
797+
798+
_PROPERTY_TO_API_FIELD = {
799+
"container_memory": "containerMemory",
800+
"container_cpu": "containerCpu",
801+
"runtime_connection": "runtimeConnection",
802+
"max_batching_rows": "maxBatchingRows",
803+
"runtime_version": "runtimeVersion",
804+
}
805+
806+
def __init__(
807+
self,
808+
container_memory: Optional[str] = None,
809+
container_cpu: Optional[int] = None,
810+
runtime_connection: Optional[str] = None,
811+
max_batching_rows: Optional[int] = None,
812+
runtime_version: Optional[str] = None,
813+
_properties: Optional[Dict] = None,
814+
) -> None:
815+
if _properties is None:
816+
_properties = {}
817+
self._properties = _properties
818+
819+
if container_memory is not None:
820+
self.container_memory = container_memory
821+
if container_cpu is not None:
822+
self.container_cpu = container_cpu
823+
if runtime_connection is not None:
824+
self.runtime_connection = runtime_connection
825+
if max_batching_rows is not None:
826+
self.max_batching_rows = max_batching_rows
827+
if runtime_version is not None:
828+
self.runtime_version = runtime_version
829+
830+
@property
831+
def container_memory(self) -> Optional[str]:
832+
"""Optional. Amount of memory provisioned for a Python UDF container instance."""
833+
return self._properties.get("containerMemory")
834+
835+
@container_memory.setter
836+
def container_memory(self, value: Optional[str]):
837+
if value is not None and not isinstance(value, str):
838+
raise ValueError("container_memory must be a string or None.")
839+
self._properties["containerMemory"] = value
840+
841+
@property
842+
def container_cpu(self) -> Optional[int]:
843+
"""Optional. Amount of CPU provisioned for a Python UDF container instance."""
844+
return self._properties.get("containerCpu")
845+
846+
@container_cpu.setter
847+
def container_cpu(self, value: Optional[int]):
848+
if value is not None and not isinstance(value, int):
849+
raise ValueError("container_cpu must be an integer or None.")
850+
self._properties["containerCpu"] = value
851+
852+
@property
853+
def runtime_connection(self) -> Optional[str]:
854+
"""Optional. Fully qualified name of the connection."""
855+
return self._properties.get("runtimeConnection")
856+
857+
@runtime_connection.setter
858+
def runtime_connection(self, value: Optional[str]):
859+
if value is not None and not isinstance(value, str):
860+
raise ValueError("runtime_connection must be a string or None.")
861+
self._properties["runtimeConnection"] = value
862+
863+
@property
864+
def max_batching_rows(self) -> Optional[int]:
865+
"""Optional. Maximum number of rows in each batch sent to the external runtime."""
866+
return _helpers._int_or_none(self._properties.get("maxBatchingRows"))
867+
868+
@max_batching_rows.setter
869+
def max_batching_rows(self, value: Optional[int]):
870+
if value is not None and not isinstance(value, int):
871+
raise ValueError("max_batching_rows must be an integer or None.")
872+
self._properties["maxBatchingRows"] = _helpers._str_or_none(value)
873+
874+
@property
875+
def runtime_version(self) -> Optional[str]:
876+
"""Optional. Language runtime version."""
877+
return self._properties.get("runtimeVersion")
878+
879+
@runtime_version.setter
880+
def runtime_version(self, value: Optional[str]):
881+
if value is not None and not isinstance(value, str):
882+
raise ValueError("runtime_version must be a string or None.")
883+
self._properties["runtimeVersion"] = value
884+
885+
@classmethod
886+
def from_api_repr(cls, resource: dict) -> "ExternalRuntimeOptions":
887+
"""Factory: construct external runtime options given its API representation.
888+
Args:
889+
resource (Dict[str, object]): Resource, as returned from the API.
890+
Returns:
891+
google.cloud.bigquery.routine.ExternalRuntimeOptions:
892+
Python object, as parsed from ``resource``.
893+
"""
894+
ref = cls()
895+
ref._properties = resource
896+
return ref
897+
898+
def to_api_repr(self) -> dict:
899+
"""Construct the API resource representation of this ExternalRuntimeOptions.
900+
Returns:
901+
Dict[str, object]: External runtime options represented as an API resource.
902+
"""
903+
return self._properties
904+
905+
def __eq__(self, other):
906+
if not isinstance(other, ExternalRuntimeOptions):
907+
return NotImplemented
908+
return self._properties == other._properties
909+
910+
def __ne__(self, other):
911+
return not self == other
912+
913+
def __repr__(self):
914+
all_properties = [
915+
"{}={}".format(property_name, repr(getattr(self, property_name)))
916+
for property_name in sorted(self._PROPERTY_TO_API_FIELD)
917+
]
918+
return "ExternalRuntimeOptions({})".format(", ".join(all_properties))
Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
# -*- coding: utf-8 -*-
2+
#
3+
# Copyright 2024 Google LLC
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# https://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
17+
import pytest
18+
19+
20+
@pytest.fixture
21+
def target_class():
22+
from google.cloud.bigquery.routine.routine import ExternalRuntimeOptions
23+
24+
return ExternalRuntimeOptions
25+
26+
27+
@pytest.fixture
28+
def object_under_test(target_class):
29+
return target_class()
30+
31+
32+
def test_ctor(target_class):
33+
container_memory = "1G"
34+
container_cpu = 1
35+
runtime_connection = "projects/my-project/locations/us-central1/connections/my-connection"
36+
max_batching_rows = 100
37+
runtime_version = "python-3.11"
38+
39+
instance = target_class(
40+
container_memory=container_memory,
41+
container_cpu=container_cpu,
42+
runtime_connection=runtime_connection,
43+
max_batching_rows=max_batching_rows,
44+
runtime_version=runtime_version,
45+
)
46+
47+
assert instance.container_memory == container_memory
48+
assert instance.container_cpu == container_cpu
49+
assert instance.runtime_connection == runtime_connection
50+
assert instance.max_batching_rows == max_batching_rows
51+
assert instance.runtime_version == runtime_version
52+
53+
54+
def test_container_memory(object_under_test):
55+
container_memory = "512Mi"
56+
object_under_test.container_memory = container_memory
57+
assert object_under_test.container_memory == container_memory
58+
59+
60+
def test_container_cpu(object_under_test):
61+
container_cpu = 1
62+
object_under_test.container_cpu = container_cpu
63+
assert object_under_test.container_cpu == container_cpu
64+
65+
66+
def test_runtime_connection(object_under_test):
67+
runtime_connection = "projects/my-project/locations/us-central1/connections/my-connection"
68+
object_under_test.runtime_connection = runtime_connection
69+
assert object_under_test.runtime_connection == runtime_connection
70+
71+
72+
def test_max_batching_rows(object_under_test):
73+
max_batching_rows = 100
74+
object_under_test.max_batching_rows = max_batching_rows
75+
assert object_under_test.max_batching_rows == max_batching_rows
76+
77+
78+
def test_runtime_version(object_under_test):
79+
runtime_version = "python-3.11"
80+
object_under_test.runtime_version = runtime_version
81+
assert object_under_test.runtime_version == runtime_version
82+
83+
84+
def test_from_api_repr(target_class):
85+
resource = {
86+
"containerMemory": "1G",
87+
"containerCpu": 1,
88+
"runtimeConnection": "projects/my-project/locations/us-central1/connections/my-connection",
89+
"maxBatchingRows": "100",
90+
"runtimeVersion": "python-3.11",
91+
}
92+
instance = target_class.from_api_repr(resource)
93+
94+
assert instance.container_memory == "1G"
95+
assert instance.container_cpu == 1
96+
assert (
97+
instance.runtime_connection
98+
== "projects/my-project/locations/us-central1/connections/my-connection"
99+
)
100+
assert instance.max_batching_rows == 100
101+
assert instance.runtime_version == "python-3.11"
102+
103+
104+
def test_to_api_repr(target_class):
105+
instance = target_class(
106+
container_memory="1G",
107+
container_cpu=1,
108+
runtime_connection="projects/my-project/locations/us-central1/connections/my-connection",
109+
max_batching_rows=100,
110+
runtime_version="python-3.11",
111+
)
112+
resource = instance.to_api_repr()
113+
114+
assert resource == {
115+
"containerMemory": "1G",
116+
"containerCpu": 1,
117+
"runtimeConnection": "projects/my-project/locations/us-central1/connections/my-connection",
118+
"maxBatchingRows": "100",
119+
"runtimeVersion": "python-3.11",
120+
}
121+
122+
123+
def test_repr(target_class):
124+
instance = target_class(
125+
container_memory="1G",
126+
container_cpu=1,
127+
)
128+
expected_repr = (
129+
"ExternalRuntimeOptions(container_cpu=1, container_memory='1G', "
130+
"max_batching_rows=None, runtime_connection=None, runtime_version=None)"
131+
)
132+
assert repr(instance) == expected_repr
133+
134+
135+
def test_invalid_container_memory(object_under_test):
136+
with pytest.raises(ValueError, match="container_memory must be a string or None."):
137+
object_under_test.container_memory = 123
138+
139+
140+
def test_invalid_container_cpu(object_under_test):
141+
with pytest.raises(ValueError, match="container_cpu must be an integer or None."):
142+
object_under_test.container_cpu = "1"
143+
144+
145+
def test_invalid_runtime_connection(object_under_test):
146+
with pytest.raises(ValueError, match="runtime_connection must be a string or None."):
147+
object_under_test.runtime_connection = 123
148+
149+
150+
def test_invalid_max_batching_rows(object_under_test):
151+
with pytest.raises(ValueError, match="max_batching_rows must be an integer or None."):
152+
object_under_test.max_batching_rows = "100"
153+
154+
155+
def test_invalid_runtime_version(object_under_test):
156+
with pytest.raises(ValueError, match="runtime_version must be a string or None."):
157+
object_under_test.runtime_version = 123

0 commit comments

Comments
 (0)