Skip to content

Commit b736b23

Browse files
chalmerlowechalmerlowe
committed
feat: adds centralized client and tests (#2249)
1 parent 698f401 commit b736b23

File tree

4 files changed

+459
-0
lines changed

4 files changed

+459
-0
lines changed

‎google/cloud/bigquery_v2/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from .services.routine_service import RoutineServiceClient
2525
from .services.row_access_policy_service import RowAccessPolicyServiceClient
2626
from .services.table_service import TableServiceClient
27+
from .services.centralized_services import BigQueryClient
2728

2829
from .types.biglake_config import BigLakeConfiguration
2930
from .types.clustering import Clustering
@@ -214,6 +215,7 @@
214215
"BiEngineReason",
215216
"BiEngineStatistics",
216217
"BigLakeConfiguration",
218+
"BigQueryClient",
217219
"BigtableColumn",
218220
"BigtableColumnFamily",
219221
"BigtableOptions",

‎google/cloud/bigquery_v2/services/centralized_services/__init__.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2025 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
from .client import BigQueryClient
17+
18+
__all__ = ("BigQueryClient",)

‎google/cloud/bigquery_v2/services/centralized_services/client.py‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2025 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
#
16+
17+
import os
18+
from typing import (
19+
Optional,
20+
Sequence,
21+
Tuple,
22+
Union,
23+
)
24+
25+
# Import service clients
26+
from google.cloud.bigquery_v2.services.dataset_service import DatasetServiceClient
27+
from google.cloud.bigquery_v2.services.job_service import JobServiceClient
28+
from google.cloud.bigquery_v2.services.model_service import ModelServiceClient
29+
30+
# Import Request classes
31+
from google.cloud.bigquery_v2.types import (
32+
# DatasetService Request classes
33+
GetDatasetRequest,
34+
# JobService Request classes
35+
ListJobsRequest,
36+
# ModelService Request classes
37+
DeleteModelRequest,
38+
GetModelRequest,
39+
PatchModelRequest,
40+
ListModelsRequest,
41+
)
42+
43+
from google.api_core import gapic_v1
44+
from google.api_core import retry as retries
45+
46+
# Create a type alias
47+
try:
48+
OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None]
49+
except AttributeError: # pragma: NO COVER
50+
OptionalRetry = Union[retries.Retry, object, None] # type: ignore
51+
52+
# TODO: revise this universally.
53+
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
54+
55+
DEFAULT_RETRY: OptionalRetry = gapic_v1.method.DEFAULT
56+
DEFAULT_TIMEOUT: Union[float, object] = gapic_v1.method.DEFAULT
57+
DEFAULT_METADATA: Sequence[Tuple[str, Union[str, bytes]]] = ()
58+
59+
60+
def _drop_self_key(kwargs):
61+
"Drops 'self' key from a given kwargs dict."
62+
63+
if not isinstance(kwargs, dict):
64+
raise TypeError("kwargs must be a dict.")
65+
kwargs.pop("self", None) # Essentially a no-op if 'self' key does not exist
66+
return kwargs
67+
68+
69+
# Create Centralized Client
70+
class BigQueryClient:
71+
def __init__(self):
72+
# Dataset service related init attributes
73+
self.dataset_service_client = DatasetServiceClient()
74+
self.job_service_client = JobServiceClient()
75+
self.model_service_client = ModelServiceClient()
76+
77+
def get_dataset(
78+
self,
79+
request: Optional[Union[GetDatasetRequest, dict]] = None,
80+
*,
81+
retry: OptionalRetry = DEFAULT_RETRY,
82+
timeout: Union[float, object] = DEFAULT_TIMEOUT,
83+
metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
84+
):
85+
"""
86+
TODO: Add docstring.
87+
"""
88+
kwargs = _drop_self_key(locals())
89+
return self.dataset_service_client.get_dataset(**kwargs)
90+
91+
def list_jobs(
92+
self,
93+
request: Optional[Union[ListJobsRequest, dict]] = None,
94+
*,
95+
retry: OptionalRetry = DEFAULT_RETRY,
96+
timeout: Union[float, object] = DEFAULT_TIMEOUT,
97+
metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
98+
):
99+
"""
100+
TODO: Add docstring.
101+
"""
102+
kwargs = _drop_self_key(locals())
103+
return self.job_service_client.list_jobs(**kwargs)
104+
105+
def get_model(
106+
self,
107+
request: Optional[Union[GetModelRequest, dict]] = None,
108+
*,
109+
retry: OptionalRetry = DEFAULT_RETRY,
110+
timeout: Union[float, object] = DEFAULT_TIMEOUT,
111+
metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
112+
):
113+
"""
114+
TODO: Add docstring.
115+
"""
116+
kwargs = _drop_self_key(locals())
117+
return self.model_service_client.get_model(**kwargs)
118+
119+
def delete_model(
120+
self,
121+
request: Optional[Union[DeleteModelRequest, dict]] = None,
122+
*,
123+
retry: OptionalRetry = DEFAULT_RETRY,
124+
timeout: Union[float, object] = DEFAULT_TIMEOUT,
125+
metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
126+
):
127+
"""
128+
TODO: Add docstring.
129+
"""
130+
kwargs = _drop_self_key(locals())
131+
# The underlying GAPIC client returns None on success.
132+
return self.model_service_client.delete_model(**kwargs)
133+
134+
def patch_model(
135+
self,
136+
request: Optional[Union[PatchModelRequest, dict]] = None,
137+
*,
138+
retry: OptionalRetry = DEFAULT_RETRY,
139+
timeout: Union[float, object] = DEFAULT_TIMEOUT,
140+
metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
141+
):
142+
"""
143+
TODO: Add docstring.
144+
"""
145+
kwargs = _drop_self_key(locals())
146+
return self.model_service_client.patch_model(**kwargs)
147+
148+
def list_models(
149+
self,
150+
request: Optional[Union[ListModelsRequest, dict]] = None,
151+
*,
152+
retry: OptionalRetry = DEFAULT_RETRY,
153+
timeout: Union[float, object] = DEFAULT_TIMEOUT,
154+
metadata: Sequence[Tuple[str, Union[str, bytes]]] = DEFAULT_METADATA,
155+
):
156+
"""
157+
TODO: Add docstring.
158+
"""
159+
kwargs = _drop_self_key(locals())
160+
return self.model_service_client.list_models(**kwargs)
161+
162+
163+
# ===============================================
164+
# Sample TODO: Relocate this to a samples file
165+
# ===============================================
166+
167+
# Instantiate BQClient class
168+
bqclient = BigQueryClient()
169+
170+
# Instantiate Request class
171+
get_dataset_request = GetDatasetRequest(
172+
project_id=PROJECT_ID,
173+
dataset_id="experimental",
174+
)
175+
176+
# Generate response
177+
dataset = bqclient.get_dataset(request=get_dataset_request)
178+
179+
# Display response
180+
print(f"GET DATASET:\n\t{dataset.id=}\n")

0 commit comments

Comments
 (0)