Skip to content
This repository was archived by the owner on Jul 6, 2023. It is now read-only.

Commit 87911bc

Browse files
feat: add context manager support in client (#51)
- [ ] 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: https://github.com/googleapis/googleapis-gen/commit/81decffe9fc72396a8153e756d1d67a6eecfd620 Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9
1 parent 904c829 commit 87911bc

File tree

7 files changed

+95
-4
lines changed

7 files changed

+95
-4
lines changed

google/cloud/essential_contacts_v1/services/essential_contacts_service/async_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -682,6 +682,12 @@ async def send_test_message(
682682
request, retry=retry, timeout=timeout, metadata=metadata,
683683
)
684684

685+
async def __aenter__(self):
686+
return self
687+
688+
async def __aexit__(self, exc_type, exc, tb):
689+
await self.transport.close()
690+
685691

686692
try:
687693
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/essential_contacts_v1/services/essential_contacts_service/client.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -348,10 +348,7 @@ def __init__(
348348
client_cert_source_for_mtls=client_cert_source_func,
349349
quota_project_id=client_options.quota_project_id,
350350
client_info=client_info,
351-
always_use_jwt_access=(
352-
Transport == type(self).get_transport_class("grpc")
353-
or Transport == type(self).get_transport_class("grpc_asyncio")
354-
),
351+
always_use_jwt_access=True,
355352
)
356353

357354
def create_contact(
@@ -849,6 +846,19 @@ def send_test_message(
849846
request, retry=retry, timeout=timeout, metadata=metadata,
850847
)
851848

849+
def __enter__(self):
850+
return self
851+
852+
def __exit__(self, type, value, traceback):
853+
"""Releases underlying transport's resources.
854+
855+
.. warning::
856+
ONLY use as a context manager if the transport is NOT shared
857+
with other clients! Exiting the with block will CLOSE the transport
858+
and may cause errors in other clients!
859+
"""
860+
self.transport.close()
861+
852862

853863
try:
854864
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(

google/cloud/essential_contacts_v1/services/essential_contacts_service/transports/base.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,15 @@ def _prep_wrapped_messages(self, client_info):
200200
),
201201
}
202202

203+
def close(self):
204+
"""Closes resources associated with the transport.
205+
206+
.. warning::
207+
Only call this method if the transport is NOT shared
208+
with other clients - this may cause errors in other clients!
209+
"""
210+
raise NotImplementedError()
211+
203212
@property
204213
def create_contact(
205214
self,

google/cloud/essential_contacts_v1/services/essential_contacts_service/transports/grpc.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,5 +410,8 @@ def send_test_message(
410410
)
411411
return self._stubs["send_test_message"]
412412

413+
def close(self):
414+
self.grpc_channel.close()
415+
413416

414417
__all__ = ("EssentialContactsServiceGrpcTransport",)

google/cloud/essential_contacts_v1/services/essential_contacts_service/transports/grpc_asyncio.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -419,5 +419,8 @@ def send_test_message(
419419
)
420420
return self._stubs["send_test_message"]
421421

422+
def close(self):
423+
return self.grpc_channel.close()
424+
422425

423426
__all__ = ("EssentialContactsServiceGrpcAsyncIOTransport",)

google/cloud/essential_contacts_v1/types/service.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939

4040
class Contact(proto.Message):
4141
r"""A contact that will receive notifications from Google Cloud.
42+
4243
Attributes:
4344
name (str):
4445
The identifier for the contact. Format:
@@ -79,6 +80,7 @@ class Contact(proto.Message):
7980

8081
class ListContactsRequest(proto.Message):
8182
r"""Request message for the ListContacts method.
83+
8284
Attributes:
8385
parent (str):
8486
Required. The parent resource name. Format:
@@ -105,6 +107,7 @@ class ListContactsRequest(proto.Message):
105107

106108
class ListContactsResponse(proto.Message):
107109
r"""Response message for the ListContacts method.
110+
108111
Attributes:
109112
contacts (Sequence[google.cloud.essential_contacts_v1.types.Contact]):
110113
The contacts for the specified resource.
@@ -126,6 +129,7 @@ def raw_page(self):
126129

127130
class GetContactRequest(proto.Message):
128131
r"""Request message for the GetContact method.
132+
129133
Attributes:
130134
name (str):
131135
Required. The name of the contact to retrieve. Format:
@@ -139,6 +143,7 @@ class GetContactRequest(proto.Message):
139143

140144
class DeleteContactRequest(proto.Message):
141145
r"""Request message for the DeleteContact method.
146+
142147
Attributes:
143148
name (str):
144149
Required. The name of the contact to delete. Format:
@@ -152,6 +157,7 @@ class DeleteContactRequest(proto.Message):
152157

153158
class CreateContactRequest(proto.Message):
154159
r"""Request message for the CreateContact method.
160+
155161
Attributes:
156162
parent (str):
157163
Required. The resource to save this contact for. Format:
@@ -168,6 +174,7 @@ class CreateContactRequest(proto.Message):
168174

169175
class UpdateContactRequest(proto.Message):
170176
r"""Request message for the UpdateContact method.
177+
171178
Attributes:
172179
contact (google.cloud.essential_contacts_v1.types.Contact):
173180
Required. The contact resource to replace the
@@ -187,6 +194,7 @@ class UpdateContactRequest(proto.Message):
187194

188195
class ComputeContactsRequest(proto.Message):
189196
r"""Request message for the ComputeContacts method.
197+
190198
Attributes:
191199
parent (str):
192200
Required. The name of the resource to compute contacts for.
@@ -221,6 +229,7 @@ class ComputeContactsRequest(proto.Message):
221229

222230
class ComputeContactsResponse(proto.Message):
223231
r"""Response message for the ComputeContacts method.
232+
224233
Attributes:
225234
contacts (Sequence[google.cloud.essential_contacts_v1.types.Contact]):
226235
All contacts for the resource that are
@@ -245,6 +254,7 @@ def raw_page(self):
245254

246255
class SendTestMessageRequest(proto.Message):
247256
r"""Request message for the SendTestMessage method.
257+
248258
Attributes:
249259
contacts (Sequence[str]):
250260
Required. The list of names of the contacts to send a test

tests/unit/gapic/essential_contacts_v1/test_essential_contacts_service.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from google.api_core import gapic_v1
3030
from google.api_core import grpc_helpers
3131
from google.api_core import grpc_helpers_async
32+
from google.api_core import path_template
3233
from google.auth import credentials as ga_credentials
3334
from google.auth.exceptions import MutualTLSChannelError
3435
from google.cloud.essential_contacts_v1.services.essential_contacts_service import (
@@ -2303,6 +2304,9 @@ def test_essential_contacts_service_base_transport():
23032304
with pytest.raises(NotImplementedError):
23042305
getattr(transport, method)(request=object())
23052306

2307+
with pytest.raises(NotImplementedError):
2308+
transport.close()
2309+
23062310

23072311
@requires_google_auth_gte_1_25_0
23082312
def test_essential_contacts_service_base_transport_with_credentials_file():
@@ -2787,3 +2791,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
27872791
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
27882792
)
27892793
prep.assert_called_once_with(client_info)
2794+
2795+
2796+
@pytest.mark.asyncio
2797+
async def test_transport_close_async():
2798+
client = EssentialContactsServiceAsyncClient(
2799+
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
2800+
)
2801+
with mock.patch.object(
2802+
type(getattr(client.transport, "grpc_channel")), "close"
2803+
) as close:
2804+
async with client:
2805+
close.assert_not_called()
2806+
close.assert_called_once()
2807+
2808+
2809+
def test_transport_close():
2810+
transports = {
2811+
"grpc": "_grpc_channel",
2812+
}
2813+
2814+
for transport, close_name in transports.items():
2815+
client = EssentialContactsServiceClient(
2816+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
2817+
)
2818+
with mock.patch.object(
2819+
type(getattr(client.transport, close_name)), "close"
2820+
) as close:
2821+
with client:
2822+
close.assert_not_called()
2823+
close.assert_called_once()
2824+
2825+
2826+
def test_client_ctx():
2827+
transports = [
2828+
"grpc",
2829+
]
2830+
for transport in transports:
2831+
client = EssentialContactsServiceClient(
2832+
credentials=ga_credentials.AnonymousCredentials(), transport=transport
2833+
)
2834+
# Test client calls underlying transport.
2835+
with mock.patch.object(type(client.transport), "close") as close:
2836+
close.assert_not_called()
2837+
with client:
2838+
pass
2839+
close.assert_called()

0 commit comments

Comments
 (0)