Skip to content
Permalink
Browse files
feat: add context manager support in client (#50)
- [ ] 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
  • Loading branch information
gcf-owl-bot[bot] committed Oct 8, 2021
1 parent 32b44f7 commit 01d2652f1a619688fe8baf7478540787cdcf7530
Show file tree
Hide file tree
Showing 8 changed files with 89 additions and 7 deletions.
@@ -1494,6 +1494,12 @@ async def delete_connection_profile(
# Done; return the response.
return response

async def __aenter__(self):
return self

async def __aexit__(self, exc_type, exc, tb):
await self.transport.close()


try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
@@ -372,10 +372,7 @@ def __init__(
client_cert_source_for_mtls=client_cert_source_func,
quota_project_id=client_options.quota_project_id,
client_info=client_info,
always_use_jwt_access=(
Transport == type(self).get_transport_class("grpc")
or Transport == type(self).get_transport_class("grpc_asyncio")
),
always_use_jwt_access=True,
)

def list_migration_jobs(
@@ -1707,6 +1704,19 @@ def delete_connection_profile(
# Done; return the response.
return response

def __enter__(self):
return self

def __exit__(self, type, value, traceback):
"""Releases underlying transport's resources.
.. warning::
ONLY use as a context manager if the transport is NOT shared
with other clients! Exiting the with block will CLOSE the transport
and may cause errors in other clients!
"""
self.transport.close()


try:
DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo(
@@ -232,6 +232,15 @@ def _prep_wrapped_messages(self, client_info):
),
}

def close(self):
"""Closes resources associated with the transport.
.. warning::
Only call this method if the transport is NOT shared
with other clients - this may cause errors in other clients!
"""
raise NotImplementedError()

@property
def operations_client(self) -> operations_v1.OperationsClient:
"""Return the client designed to process long-running operations."""
@@ -707,5 +707,8 @@ def delete_connection_profile(
)
return self._stubs["delete_connection_profile"]

def close(self):
self.grpc_channel.close()


__all__ = ("DataMigrationServiceGrpcTransport",)
@@ -740,5 +740,8 @@ def delete_connection_profile(
)
return self._stubs["delete_connection_profile"]

def close(self):
return self.grpc_channel.close()


__all__ = ("DataMigrationServiceGrpcAsyncIOTransport",)
@@ -329,7 +329,7 @@ class GenerateSshScriptRequest(proto.Message):


class VmCreationConfig(proto.Message):
r"""VM creation configuration messag.
r"""VM creation configuration message
Attributes:
vm_machine_type (str):
@@ -348,7 +348,7 @@ class VmCreationConfig(proto.Message):


class VmSelectionConfig(proto.Message):
r"""VM selection configuration messag.
r"""VM selection configuration message
Attributes:
vm_zone (str):
@@ -395,7 +395,8 @@ class StaticIpConnectivity(proto.Message):
destination database's public IP. You can retrieve the Cloud SQL
instance's public IP from the Cloud SQL console or using Cloud
SQL APIs. No additional configuration is required.
"""
"""


class ReverseSshConnectivity(proto.Message):
@@ -32,6 +32,7 @@
from google.api_core import grpc_helpers_async
from google.api_core import operation_async # type: ignore
from google.api_core import operations_v1
from google.api_core import path_template
from google.auth import credentials as ga_credentials
from google.auth.exceptions import MutualTLSChannelError
from google.cloud.clouddms_v1.services.data_migration_service import (
@@ -4388,6 +4389,9 @@ def test_data_migration_service_base_transport():
with pytest.raises(NotImplementedError):
getattr(transport, method)(request=object())

with pytest.raises(NotImplementedError):
transport.close()

# Additionally, the LRO client (a property) should
# also raise NotImplementedError
with pytest.raises(NotImplementedError):
@@ -4931,3 +4935,49 @@ def test_client_withDEFAULT_CLIENT_INFO():
credentials=ga_credentials.AnonymousCredentials(), client_info=client_info,
)
prep.assert_called_once_with(client_info)


@pytest.mark.asyncio
async def test_transport_close_async():
client = DataMigrationServiceAsyncClient(
credentials=ga_credentials.AnonymousCredentials(), transport="grpc_asyncio",
)
with mock.patch.object(
type(getattr(client.transport, "grpc_channel")), "close"
) as close:
async with client:
close.assert_not_called()
close.assert_called_once()


def test_transport_close():
transports = {
"grpc": "_grpc_channel",
}

for transport, close_name in transports.items():
client = DataMigrationServiceClient(
credentials=ga_credentials.AnonymousCredentials(), transport=transport
)
with mock.patch.object(
type(getattr(client.transport, close_name)), "close"
) as close:
with client:
close.assert_not_called()
close.assert_called_once()


def test_client_ctx():
transports = [
"grpc",
]
for transport in transports:
client = DataMigrationServiceClient(
credentials=ga_credentials.AnonymousCredentials(), transport=transport
)
# Test client calls underlying transport.
with mock.patch.object(type(client.transport), "close") as close:
close.assert_not_called()
with client:
pass
close.assert_called()

0 comments on commit 01d2652

Please sign in to comment.