Skip to content
This repository has been archived by the owner on Sep 5, 2023. It is now read-only.

feat: add context manager support in client #119

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Expand Up @@ -1094,6 +1094,12 @@ async def test_iam_permissions(
# 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(
Expand Down
Expand Up @@ -356,10 +356,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_functions(
Expand Down Expand Up @@ -1240,6 +1237,19 @@ def test_iam_permissions(
# 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(
Expand Down
Expand Up @@ -243,6 +243,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."""
Expand Down
Expand Up @@ -586,5 +586,8 @@ def test_iam_permissions(
)
return self._stubs["test_iam_permissions"]

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


__all__ = ("CloudFunctionsServiceGrpcTransport",)
Expand Up @@ -603,5 +603,8 @@ def test_iam_permissions(
)
return self._stubs["test_iam_permissions"]

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


__all__ = ("CloudFunctionsServiceGrpcAsyncIOTransport",)
15 changes: 14 additions & 1 deletion google/cloud/functions_v1/types/functions.py
Expand Up @@ -363,13 +363,15 @@ class Retry(proto.Message):
failed execution will be retried up to 7 days with an
exponential backoff (capped at 10 seconds).
Retried execution is charged as any other execution.
"""

"""

retry = proto.Field(proto.MESSAGE, number=1, oneof="action", message=Retry,)


class CreateFunctionRequest(proto.Message):
r"""Request for the ``CreateFunction`` method.

Attributes:
location (str):
Required. The project and location in which the function
Expand All @@ -385,6 +387,7 @@ class CreateFunctionRequest(proto.Message):

class UpdateFunctionRequest(proto.Message):
r"""Request for the ``UpdateFunction`` method.

Attributes:
function (google.cloud.functions_v1.types.CloudFunction):
Required. New version of the function.
Expand All @@ -401,6 +404,7 @@ class UpdateFunctionRequest(proto.Message):

class GetFunctionRequest(proto.Message):
r"""Request for the ``GetFunction`` method.

Attributes:
name (str):
Required. The name of the function which
Expand All @@ -412,6 +416,7 @@ class GetFunctionRequest(proto.Message):

class ListFunctionsRequest(proto.Message):
r"""Request for the ``ListFunctions`` method.

Attributes:
parent (str):
The project and location from which the function should be
Expand All @@ -438,6 +443,7 @@ class ListFunctionsRequest(proto.Message):

class ListFunctionsResponse(proto.Message):
r"""Response for the ``ListFunctions`` method.

Attributes:
functions (Sequence[google.cloud.functions_v1.types.CloudFunction]):
The functions that match the request.
Expand All @@ -463,6 +469,7 @@ def raw_page(self):

class DeleteFunctionRequest(proto.Message):
r"""Request for the ``DeleteFunction`` method.

Attributes:
name (str):
Required. The name of the function which
Expand All @@ -474,6 +481,7 @@ class DeleteFunctionRequest(proto.Message):

class CallFunctionRequest(proto.Message):
r"""Request for the ``CallFunction`` method.

Attributes:
name (str):
Required. The name of the function to be
Expand All @@ -488,6 +496,7 @@ class CallFunctionRequest(proto.Message):

class CallFunctionResponse(proto.Message):
r"""Response of ``CallFunction`` method.

Attributes:
execution_id (str):
Execution id of function invocation.
Expand All @@ -508,6 +517,7 @@ class CallFunctionResponse(proto.Message):

class GenerateUploadUrlRequest(proto.Message):
r"""Request of ``GenerateSourceUploadUrl`` method.

Attributes:
parent (str):
The project and location in which the Google Cloud Storage
Expand All @@ -520,6 +530,7 @@ class GenerateUploadUrlRequest(proto.Message):

class GenerateUploadUrlResponse(proto.Message):
r"""Response of ``GenerateSourceUploadUrl`` method.

Attributes:
upload_url (str):
The generated Google Cloud Storage signed URL
Expand All @@ -533,6 +544,7 @@ class GenerateUploadUrlResponse(proto.Message):

class GenerateDownloadUrlRequest(proto.Message):
r"""Request of ``GenerateDownloadUrl`` method.

Attributes:
name (str):
The name of function for which source code
Expand All @@ -549,6 +561,7 @@ class GenerateDownloadUrlRequest(proto.Message):

class GenerateDownloadUrlResponse(proto.Message):
r"""Response of ``GenerateDownloadUrl`` method.

Attributes:
download_url (str):
The generated Google Cloud Storage signed URL
Expand Down
1 change: 1 addition & 0 deletions google/cloud/functions_v1/types/operations.py
Expand Up @@ -35,6 +35,7 @@ class OperationType(proto.Enum):

class OperationMetadataV1(proto.Message):
r"""Metadata describing an [Operation][google.longrunning.Operation]

Attributes:
target (str):
Target of the operation - for example
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/gapic/functions_v1/test_cloud_functions_service.py
Expand Up @@ -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.functions_v1.services.cloud_functions_service import (
Expand Down Expand Up @@ -2847,6 +2848,9 @@ def test_cloud_functions_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):
Expand Down Expand Up @@ -3366,3 +3370,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 = CloudFunctionsServiceAsyncClient(
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 = CloudFunctionsServiceClient(
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 = CloudFunctionsServiceClient(
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()