Skip to content

Commit 17bb546

Browse files
feat: add context manager support in client (#87)
- [ ] Regenerate this pull request now. chore: fix docstring for first attribute of protos committer: @busunkim96 PiperOrigin-RevId: 401271153 Source-Link: googleapis/googleapis@787f8c9 Source-Link: googleapis/googleapis-gen@81decff Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9
1 parent d342862 commit 17bb546

File tree

8 files changed

+97
-16
lines changed

8 files changed

+97
-16
lines changed

packages/google-cloud-gke-hub/google/cloud/gkehub_v1/configmanagement_v1/types/configmanagement.py

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ class MembershipSpec(proto.Message):
123123

124124

125125
class ConfigSync(proto.Message):
126-
r"""Configuration for Config Syn.
126+
r"""Configuration for Config Sync
127127
128128
Attributes:
129129
git (google.cloud.gkehub.configmanagement_v1.types.GitConfig):
@@ -179,7 +179,7 @@ class GitConfig(proto.Message):
179179

180180

181181
class PolicyController(proto.Message):
182-
r"""Configuration for Policy Controlle.
182+
r"""Configuration for Policy Controller
183183
184184
Attributes:
185185
enabled (bool):
@@ -214,7 +214,7 @@ class PolicyController(proto.Message):
214214

215215

216216
class HierarchyControllerConfig(proto.Message):
217-
r"""Configuration for Hierarchy Controlle.
217+
r"""Configuration for Hierarchy Controller
218218
219219
Attributes:
220220
enabled (bool):
@@ -234,7 +234,7 @@ class HierarchyControllerConfig(proto.Message):
234234

235235

236236
class HierarchyControllerDeploymentState(proto.Message):
237-
r"""Deployment state for Hierarchy Controlle.
237+
r"""Deployment state for Hierarchy Controller
238238
239239
Attributes:
240240
hnc (google.cloud.gkehub.configmanagement_v1.types.DeploymentState):
@@ -250,7 +250,7 @@ class HierarchyControllerDeploymentState(proto.Message):
250250

251251

252252
class HierarchyControllerVersion(proto.Message):
253-
r"""Version for Hierarchy Controlle.
253+
r"""Version for Hierarchy Controller
254254
255255
Attributes:
256256
hnc (str):
@@ -264,7 +264,7 @@ class HierarchyControllerVersion(proto.Message):
264264

265265

266266
class HierarchyControllerState(proto.Message):
267-
r"""State for Hierarchy Controlle.
267+
r"""State for Hierarchy Controller
268268
269269
Attributes:
270270
version (google.cloud.gkehub.configmanagement_v1.types.HierarchyControllerVersion):
@@ -282,7 +282,7 @@ class HierarchyControllerState(proto.Message):
282282

283283

284284
class OperatorState(proto.Message):
285-
r"""State information for an ACM's Operato.
285+
r"""State information for an ACM's Operator
286286
287287
Attributes:
288288
version (str):
@@ -299,7 +299,7 @@ class OperatorState(proto.Message):
299299

300300

301301
class InstallError(proto.Message):
302-
r"""Errors pertaining to the installation of AC.
302+
r"""Errors pertaining to the installation of ACM
303303
304304
Attributes:
305305
error_message (str):
@@ -311,7 +311,7 @@ class InstallError(proto.Message):
311311

312312

313313
class ConfigSyncState(proto.Message):
314-
r"""State information for ConfigSyn.
314+
r"""State information for ConfigSync
315315
316316
Attributes:
317317
version (google.cloud.gkehub.configmanagement_v1.types.ConfigSyncVersion):
@@ -362,7 +362,7 @@ class ConfigSyncVersion(proto.Message):
362362

363363

364364
class ConfigSyncDeploymentState(proto.Message):
365-
r"""The state of ConfigSync's deployment on a cluste.
365+
r"""The state of ConfigSync's deployment on a cluster
366366
367367
Attributes:
368368
importer (google.cloud.gkehub.configmanagement_v1.types.DeploymentState):
@@ -484,7 +484,7 @@ class ErrorResource(proto.Message):
484484

485485

486486
class GroupVersionKind(proto.Message):
487-
r"""A Kubernetes object's GV.
487+
r"""A Kubernetes object's GVK
488488
489489
Attributes:
490490
group (str):

packages/google-cloud-gke-hub/google/cloud/gkehub_v1/services/gke_hub/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1147,6 +1147,12 @@ async def generate_connect_manifest(
11471147
# Done; return the response.
11481148
return response
11491149

1150+
async def __aenter__(self):
1151+
return self
1152+
1153+
async def __aexit__(self, exc_type, exc, tb):
1154+
await self.transport.close()
1155+
11501156

11511157
try:
11521158
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

packages/google-cloud-gke-hub/google/cloud/gkehub_v1/services/gke_hub/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -379,10 +379,7 @@ def __init__(
379379
client_cert_source_for_mtls=client_cert_source_func,
380380
quota_project_id=client_options.quota_project_id,
381381
client_info=client_info,
382-
always_use_jwt_access=(
383-
Transport == type(self).get_transport_class("grpc")
384-
or Transport == type(self).get_transport_class("grpc_asyncio")
385-
),
382+
always_use_jwt_access=True,
386383
)
387384

388385
def list_memberships(
@@ -1361,6 +1358,19 @@ def generate_connect_manifest(
13611358
# Done; return the response.
13621359
return response
13631360

1361+
def __enter__(self):
1362+
return self
1363+
1364+
def __exit__(self, type, value, traceback):
1365+
"""Releases underlying transport's resources.
1366+
1367+
.. warning::
1368+
ONLY use as a context manager if the transport is NOT shared
1369+
with other clients! Exiting the with block will CLOSE the transport
1370+
and may cause errors in other clients!
1371+
"""
1372+
self.transport.close()
1373+
13641374

13651375
try:
13661376
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

packages/google-cloud-gke-hub/google/cloud/gkehub_v1/services/gke_hub/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,15 @@ def _prep_wrapped_messages(self, client_info):
193193
),
194194
}
195195

196+
def close(self):
197+
"""Closes resources associated with the transport.
198+
199+
.. warning::
200+
Only call this method if the transport is NOT shared
201+
with other clients - this may cause errors in other clients!
202+
"""
203+
raise NotImplementedError()
204+
196205
@property
197206
def operations_client(self) -> operations_v1.OperationsClient:
198207
"""Return the client designed to process long-running operations."""

packages/google-cloud-gke-hub/google/cloud/gkehub_v1/services/gke_hub/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -556,5 +556,8 @@ def generate_connect_manifest(
556556
)
557557
return self._stubs["generate_connect_manifest"]
558558

559+
def close(self):
560+
self.grpc_channel.close()
561+
559562

560563
__all__ = ("GkeHubGrpcTransport",)

packages/google-cloud-gke-hub/google/cloud/gkehub_v1/services/gke_hub/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -573,5 +573,8 @@ def generate_connect_manifest(
573573
)
574574
return self._stubs["generate_connect_manifest"]
575575

576+
def close(self):
577+
return self.grpc_channel.close()
578+
576579

577580
__all__ = ("GkeHubGrpcAsyncIOTransport",)

packages/google-cloud-gke-hub/google/cloud/gkehub_v1/types/feature.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,7 @@ class Code(proto.Enum):
172172

173173

174174
class CommonFeatureSpec(proto.Message):
175-
r"""CommonFeatureSpec contains Hub-wide configuration informatio.
175+
r"""CommonFeatureSpec contains Hub-wide configuration information
176176
177177
Attributes:
178178
multiclusteringress (google.cloud.gkehub_v1.multiclusteringress_v1.FeatureSpec):

packages/google-cloud-gke-hub/tests/unit/gapic/gkehub_v1/test_gke_hub.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
from google.api_core import grpc_helpers_async
3333
from google.api_core import operation_async # type: ignore
3434
from google.api_core import operations_v1
35+
from google.api_core import path_template
3536
from google.auth import credentials as ga_credentials
3637
from google.auth.exceptions import MutualTLSChannelError
3738
from google.cloud.gkehub_v1 import configmanagement_v1 # type: ignore
@@ -3124,6 +3125,9 @@ def test_gke_hub_base_transport():
31243125
with pytest.raises(NotImplementedError):
31253126
getattr(transport, method)(request=object())
31263127

3128+
with pytest.raises(NotImplementedError):
3129+
transport.close()
3130+
31273131
# Additionally, the LRO client (a property) should
31283132
# also raise NotImplementedError
31293133
with pytest.raises(NotImplementedError):
@@ -3644,3 +3648,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
36443648
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
36453649
)
36463650
prep.assert_called_once_with(client_info)
3651+
3652+
3653+
@pytest.mark.asyncio
3654+
async def test_transport_close_async():
3655+
client = GkeHubAsyncClient(
3656+
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
3657+
)
3658+
with mock.patch.object(
3659+
type(getattr(client.transport, "grpc_channel")), "close"
3660+
) as close:
3661+
async with client:
3662+
close.assert_not_called()
3663+
close.assert_called_once()
3664+
3665+
3666+
def test_transport_close():
3667+
transports = {
3668+
"grpc": "_grpc_channel",
3669+
}
3670+
3671+
for transport, close_name in transports.items():
3672+
client = GkeHubClient(
3673+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
3674+
)
3675+
with mock.patch.object(
3676+
type(getattr(client.transport, close_name)), "close"
3677+
) as close:
3678+
with client:
3679+
close.assert_not_called()
3680+
close.assert_called_once()
3681+
3682+
3683+
def test_client_ctx():
3684+
transports = [
3685+
"grpc",
3686+
]
3687+
for transport in transports:
3688+
client = GkeHubClient(
3689+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
3690+
)
3691+
# Test client calls underlying transport.
3692+
with mock.patch.object(type(client.transport), "close") as close:
3693+
close.assert_not_called()
3694+
with client:
3695+
pass
3696+
close.assert_called()

0 commit comments

Comments
 (0)