Skip to content

Commit 5674587

Browse files
authored
feat(experimental): Add grpc client (#1533)
* feat: Add grpc client * add an option to configure direct path * fix lint errors * move the files to _experimental folder * move the test file under tests/unit * fix patch errors * update the imports for _storage_v2 * change the default directpath value to True * add getter for grpc_client * resolve comments * add more tests * minor change * adding helper function for credentials
1 parent dc3756d commit 5674587

File tree

2 files changed

+330
-0
lines changed

2 files changed

+330
-0
lines changed
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""A client for interacting with Google Cloud Storage using the gRPC API."""
16+
17+
from google.cloud.client import ClientWithProject
18+
from google.cloud import _storage_v2 as storage_v2
19+
20+
_marker = object()
21+
22+
23+
class GrpcClient(ClientWithProject):
24+
"""A client for interacting with Google Cloud Storage using the gRPC API.
25+
26+
:type project: str or None
27+
:param project: The project which the client acts on behalf of. If not
28+
passed, falls back to the default inferred from the
29+
environment.
30+
31+
:type credentials: :class:`~google.auth.credentials.Credentials`
32+
:param credentials: (Optional) The OAuth2 Credentials to use for this
33+
client. If not passed, falls back to the default
34+
inferred from the environment.
35+
36+
:type client_info: :class:`~google.api_core.client_info.ClientInfo`
37+
:param client_info:
38+
The client info used to send a user-agent string along with API
39+
requests. If ``None``, then default info will be used. Generally,
40+
you only need to set this if you're developing your own library
41+
or partner tool.
42+
43+
:type client_options: :class:`~google.api_core.client_options.ClientOptions` or :class:`dict`
44+
:param client_options: (Optional) Client options used to set user options
45+
on the client. A non-default universe domain or API endpoint should be
46+
set through client_options.
47+
48+
:type api_key: string
49+
:param api_key:
50+
(Optional) An API key. Mutually exclusive with any other credentials.
51+
This parameter is an alias for setting `client_options.api_key` and
52+
will supersede any API key set in the `client_options` parameter.
53+
54+
:type attempt_direct_path: bool
55+
:param attempt_direct_path:
56+
(Optional) Whether to attempt to use DirectPath for gRPC connections.
57+
This provides a direct, unproxied connection to GCS for lower latency
58+
and higher throughput, and is highly recommended when running on Google
59+
Cloud infrastructure. Defaults to ``True``.
60+
"""
61+
62+
def __init__(
63+
self,
64+
project=_marker,
65+
credentials=None,
66+
client_info=None,
67+
client_options=None,
68+
*,
69+
api_key=None,
70+
attempt_direct_path=True,
71+
):
72+
super(GrpcClient, self).__init__(project=project, credentials=credentials)
73+
74+
if isinstance(client_options, dict):
75+
if api_key:
76+
client_options["api_key"] = api_key
77+
elif client_options is None:
78+
client_options = {} if not api_key else {"api_key": api_key}
79+
elif api_key:
80+
client_options.api_key = api_key
81+
82+
self._grpc_client = self._create_gapic_client(
83+
credentials=credentials,
84+
client_info=client_info,
85+
client_options=client_options,
86+
attempt_direct_path=attempt_direct_path,
87+
)
88+
89+
def _create_gapic_client(
90+
self,
91+
credentials=None,
92+
client_info=None,
93+
client_options=None,
94+
attempt_direct_path=True,
95+
):
96+
"""Creates and configures the low-level GAPIC `storage_v2` client."""
97+
transport_cls = storage_v2.StorageClient.get_transport_class("grpc")
98+
99+
channel = transport_cls.create_channel(attempt_direct_path=attempt_direct_path)
100+
101+
transport = transport_cls(credentials=credentials, channel=channel)
102+
103+
return storage_v2.StorageClient(
104+
credentials=credentials,
105+
transport=transport,
106+
client_info=client_info,
107+
client_options=client_options,
108+
)
109+
110+
@property
111+
def grpc_client(self):
112+
"""The underlying gRPC client.
113+
114+
This property gives users direct access to the `storage_v2.StorageClient`
115+
instance. This can be useful for accessing
116+
newly added or experimental RPCs that are not yet exposed through
117+
the high-level GrpcClient.
118+
119+
Returns:
120+
google.cloud.storage_v2.StorageClient: The configured GAPIC client.
121+
"""
122+
return self._grpc_client

tests/unit/test_grpc_client.py

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import unittest
16+
from unittest import mock
17+
from google.auth import credentials as auth_credentials
18+
from google.api_core import client_options as client_options_lib
19+
20+
21+
def _make_credentials(spec=None):
22+
if spec is None:
23+
return mock.Mock(spec=auth_credentials.Credentials)
24+
return mock.Mock(spec=spec)
25+
26+
27+
class TestGrpcClient(unittest.TestCase):
28+
@mock.patch("google.cloud.client.ClientWithProject.__init__")
29+
@mock.patch("google.cloud._storage_v2.StorageClient")
30+
def test_constructor_defaults_and_options(
31+
self, mock_storage_client, mock_base_client
32+
):
33+
from google.cloud.storage._experimental import grpc_client
34+
35+
mock_transport_cls = mock.MagicMock()
36+
mock_storage_client.get_transport_class.return_value = mock_transport_cls
37+
mock_creds = _make_credentials(spec=["_base", "_get_project_id"])
38+
mock_client_info = mock.Mock()
39+
client_options_dict = {"api_endpoint": "test.endpoint"}
40+
41+
mock_base_instance = mock_base_client.return_value
42+
mock_base_instance._credentials = mock_creds
43+
44+
client = grpc_client.GrpcClient(
45+
project="test-project",
46+
credentials=mock_creds,
47+
client_info=mock_client_info,
48+
client_options=client_options_dict,
49+
)
50+
51+
# 1. Assert that the base class was initialized correctly.
52+
mock_base_client.assert_called_once_with(
53+
project="test-project", credentials=mock_creds
54+
)
55+
56+
# 2. Assert DirectPath is ON by default.
57+
mock_storage_client.get_transport_class.assert_called_once_with("grpc")
58+
mock_transport_cls.create_channel.assert_called_once_with(
59+
attempt_direct_path=True
60+
)
61+
62+
# 3. Assert the GAPIC client was created with the correct options.
63+
mock_transport = mock_transport_cls.return_value
64+
mock_storage_client.assert_called_once_with(
65+
credentials=mock_creds,
66+
transport=mock_transport,
67+
client_info=mock_client_info,
68+
client_options=client_options_dict,
69+
)
70+
71+
# 4. Assert the client instance holds the mocked GAPIC client.
72+
self.assertIs(client.grpc_client, mock_storage_client.return_value)
73+
74+
@mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
75+
@mock.patch("google.cloud._storage_v2.StorageClient")
76+
def test_constructor_disables_direct_path(
77+
self, mock_storage_client, mock_base_client
78+
):
79+
from google.cloud.storage._experimental import grpc_client
80+
81+
mock_transport_cls = mock.MagicMock()
82+
mock_storage_client.get_transport_class.return_value = mock_transport_cls
83+
mock_creds = _make_credentials()
84+
mock_base_instance = mock_base_client.return_value
85+
mock_base_instance._credentials = mock_creds
86+
87+
grpc_client.GrpcClient(
88+
project="test-project",
89+
credentials=mock_creds,
90+
attempt_direct_path=False,
91+
)
92+
93+
mock_transport_cls.create_channel.assert_called_once_with(
94+
attempt_direct_path=False
95+
)
96+
97+
@mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
98+
@mock.patch("google.cloud._storage_v2.StorageClient")
99+
def test_constructor_initialize_with_api_key(
100+
self, mock_storage_client, mock_base_client
101+
):
102+
from google.cloud.storage._experimental import grpc_client
103+
104+
mock_transport_cls = mock.MagicMock()
105+
mock_storage_client.get_transport_class.return_value = mock_transport_cls
106+
mock_creds = _make_credentials()
107+
mock_creds.project_id = None
108+
109+
mock_base_instance = mock_base_client.return_value
110+
mock_base_instance._credentials = mock_creds
111+
112+
# Instantiate with just the api_key.
113+
grpc_client.GrpcClient(
114+
project="test-project", credentials=mock_creds, api_key="test-api-key"
115+
)
116+
117+
# Assert that the GAPIC client was called with client_options
118+
# that contains the api_key.
119+
mock_transport = mock_transport_cls.return_value
120+
mock_storage_client.assert_called_once_with(
121+
credentials=mock_creds,
122+
transport=mock_transport,
123+
client_info=None,
124+
client_options={"api_key": "test-api-key"},
125+
)
126+
127+
@mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
128+
@mock.patch("google.cloud._storage_v2.StorageClient")
129+
def test_grpc_client_property(self, mock_storage_client, mock_base_client):
130+
from google.cloud.storage._experimental import grpc_client
131+
132+
mock_creds = _make_credentials()
133+
mock_base_client.return_value._credentials = mock_creds
134+
135+
client = grpc_client.GrpcClient(project="test-project", credentials=mock_creds)
136+
137+
retrieved_client = client.grpc_client
138+
139+
self.assertIs(retrieved_client, mock_storage_client.return_value)
140+
141+
@mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
142+
@mock.patch("google.cloud._storage_v2.StorageClient")
143+
def test_constructor_with_api_key_and_client_options(
144+
self, mock_storage_client, mock_base_client
145+
):
146+
from google.cloud.storage._experimental import grpc_client
147+
148+
mock_transport_cls = mock.MagicMock()
149+
mock_storage_client.get_transport_class.return_value = mock_transport_cls
150+
mock_transport = mock_transport_cls.return_value
151+
152+
mock_creds = _make_credentials()
153+
mock_base_instance = mock_base_client.return_value
154+
mock_base_instance._credentials = mock_creds
155+
156+
client_options_obj = client_options_lib.ClientOptions(
157+
api_endpoint="test.endpoint"
158+
)
159+
self.assertIsNone(client_options_obj.api_key)
160+
161+
grpc_client.GrpcClient(
162+
project="test-project",
163+
credentials=mock_creds,
164+
client_options=client_options_obj,
165+
api_key="new-test-key",
166+
)
167+
168+
mock_storage_client.assert_called_once_with(
169+
credentials=mock_creds,
170+
transport=mock_transport,
171+
client_info=None,
172+
client_options=client_options_obj,
173+
)
174+
self.assertEqual(client_options_obj.api_key, "new-test-key")
175+
176+
@mock.patch("google.cloud.storage._experimental.grpc_client.ClientWithProject")
177+
@mock.patch("google.cloud._storage_v2.StorageClient")
178+
def test_constructor_with_api_key_and_dict_options(
179+
self, mock_storage_client, mock_base_client
180+
):
181+
from google.cloud.storage._experimental import grpc_client
182+
183+
mock_creds = _make_credentials()
184+
mock_base_instance = mock_base_client.return_value
185+
mock_base_instance._credentials = mock_creds
186+
mock_transport_cls = mock.MagicMock()
187+
mock_storage_client.get_transport_class.return_value = mock_transport_cls
188+
mock_transport = mock_transport_cls.return_value
189+
190+
client_options_dict = {"api_endpoint": "test.endpoint"}
191+
192+
grpc_client.GrpcClient(
193+
project="test-project",
194+
credentials=mock_creds,
195+
client_options=client_options_dict,
196+
api_key="new-test-key",
197+
)
198+
199+
expected_options = {
200+
"api_endpoint": "test.endpoint",
201+
"api_key": "new-test-key",
202+
}
203+
mock_storage_client.assert_called_once_with(
204+
credentials=mock_creds,
205+
transport=mock_transport,
206+
client_info=None,
207+
client_options=expected_options,
208+
)

0 commit comments

Comments
 (0)