feat: add context manager support in client (#129)

- [ ] 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
googleapis · Oct 7, 2021 · 7b5d09e · 7b5d09e
1 parent a22b43a
commit 7b5d09e
Show file tree

Hide file tree

Showing 14 changed files with 194 additions and 8 deletions.
diff --git a/google/cloud/webrisk_v1/services/web_risk_service/async_client.py b/google/cloud/webrisk_v1/services/web_risk_service/async_client.py
@@ -538,6 +538,12 @@ async def create_submission(
         # 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(

diff --git a/google/cloud/webrisk_v1/services/web_risk_service/client.py b/google/cloud/webrisk_v1/services/web_risk_service/client.py
@@ -330,10 +330,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 compute_threat_list_diff(
@@ -683,6 +680,19 @@ def create_submission(
         # 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(

diff --git a/google/cloud/webrisk_v1/services/web_risk_service/transports/base.py b/google/cloud/webrisk_v1/services/web_risk_service/transports/base.py
@@ -202,6 +202,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 compute_threat_list_diff(
         self,

diff --git a/google/cloud/webrisk_v1/services/web_risk_service/transports/grpc.py b/google/cloud/webrisk_v1/services/web_risk_service/transports/grpc.py
@@ -356,5 +356,8 @@ def create_submission(
             )
         return self._stubs["create_submission"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("WebRiskServiceGrpcTransport",)
diff --git a/google/cloud/webrisk_v1/services/web_risk_service/transports/grpc_asyncio.py b/google/cloud/webrisk_v1/services/web_risk_service/transports/grpc_asyncio.py
@@ -362,5 +362,8 @@ def create_submission(
             )
         return self._stubs["create_submission"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("WebRiskServiceGrpcAsyncIOTransport",)
diff --git a/google/cloud/webrisk_v1/types/webrisk.py b/google/cloud/webrisk_v1/types/webrisk.py
@@ -59,6 +59,7 @@ class CompressionType(proto.Enum):
 
 class ComputeThreatListDiffRequest(proto.Message):
     r"""Describes an API diff request.
+
     Attributes:
         threat_type (google.cloud.webrisk_v1.types.ThreatType):
             Required. The threat list to update. Only a
@@ -78,6 +79,7 @@ class ComputeThreatListDiffRequest(proto.Message):
 
     class Constraints(proto.Message):
         r"""The constraints for this diff.
+
         Attributes:
             max_diff_entries (int):
                 The maximum size in number of entries. The diff will not
@@ -107,6 +109,7 @@ class Constraints(proto.Message):
 
 class ComputeThreatListDiffResponse(proto.Message):
     r"""
+
     Attributes:
         response_type (google.cloud.webrisk_v1.types.ComputeThreatListDiffResponse.ResponseType):
             The type of response. This may indicate that
@@ -147,6 +150,7 @@ class ResponseType(proto.Enum):
 
     class Checksum(proto.Message):
         r"""The expected state of a client's local database.
+
         Attributes:
             sha256 (bytes):
                 The SHA256 hash of the client state; that is,
@@ -168,6 +172,7 @@ class Checksum(proto.Message):
 
 class SearchUrisRequest(proto.Message):
     r"""Request to check URI entries against threatLists.
+
     Attributes:
         uri (str):
             Required. The URI to be checked for matches.
@@ -182,6 +187,7 @@ class SearchUrisRequest(proto.Message):
 
 class SearchUrisResponse(proto.Message):
     r"""
+
     Attributes:
         threat (google.cloud.webrisk_v1.types.SearchUrisResponse.ThreatUri):
             The threat list matches. This may be empty if
@@ -190,6 +196,7 @@ class SearchUrisResponse(proto.Message):
 
     class ThreatUri(proto.Message):
         r"""Contains threat information on a matching uri.
+
         Attributes:
             threat_types (Sequence[google.cloud.webrisk_v1.types.ThreatType]):
                 The ThreatList this threat belongs to.
@@ -227,6 +234,7 @@ class SearchHashesRequest(proto.Message):
 
 class SearchHashesResponse(proto.Message):
     r"""
+
     Attributes:
         threats (Sequence[google.cloud.webrisk_v1.types.SearchHashesResponse.ThreatHash]):
             The full hashes that matched the requested
@@ -239,6 +247,7 @@ class SearchHashesResponse(proto.Message):
 
     class ThreatHash(proto.Message):
         r"""Contains threat information on a matching hash.
+
         Attributes:
             threat_types (Sequence[google.cloud.webrisk_v1.types.ThreatType]):
                 The ThreatList this threat belongs to.
@@ -288,6 +297,7 @@ class ThreatEntryAdditions(proto.Message):
 
 class ThreatEntryRemovals(proto.Message):
     r"""Contains the set of entries to remove from a local database.
+
     Attributes:
         raw_indices (google.cloud.webrisk_v1.types.RawIndices):
             The raw removal indices for a local list.
@@ -305,6 +315,7 @@ class ThreatEntryRemovals(proto.Message):
 
 class RawIndices(proto.Message):
     r"""A set of raw indices to remove from a local list.
+
     Attributes:
         indices (Sequence[int]):
             The indices to remove from a
@@ -373,6 +384,7 @@ class RiceDeltaEncoding(proto.Message):
 
 class Submission(proto.Message):
     r"""Wraps a URI that might be displaying phishing content.
+
     Attributes:
         uri (str):
             Required. The URI that is being reported for
@@ -384,6 +396,7 @@ class Submission(proto.Message):
 
 class CreateSubmissionRequest(proto.Message):
     r"""Request to send a potentially phishy URI to WebRisk.
+
     Attributes:
         parent (str):
             Required. The name of the project that is making the

diff --git a/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/async_client.py b/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/async_client.py
@@ -430,6 +430,12 @@ async def search_hashes(
         # 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(

diff --git a/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/client.py b/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/client.py
@@ -332,10 +332,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 compute_threat_list_diff(
@@ -574,6 +571,19 @@ def search_hashes(
         # 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(

diff --git a/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/transports/base.py b/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/transports/base.py
@@ -199,6 +199,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 compute_threat_list_diff(
         self,

diff --git a/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/transports/grpc.py b/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/transports/grpc.py
@@ -312,5 +312,8 @@ def search_hashes(
             )
         return self._stubs["search_hashes"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("WebRiskServiceV1Beta1GrpcTransport",)
diff --git a/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/transports/grpc_asyncio.py b/google/cloud/webrisk_v1beta1/services/web_risk_service_v1_beta1/transports/grpc_asyncio.py
@@ -318,5 +318,8 @@ def search_hashes(
             )
         return self._stubs["search_hashes"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("WebRiskServiceV1Beta1GrpcAsyncIOTransport",)
diff --git a/google/cloud/webrisk_v1beta1/types/webrisk.py b/google/cloud/webrisk_v1beta1/types/webrisk.py
@@ -57,6 +57,7 @@ class CompressionType(proto.Enum):
 
 class ComputeThreatListDiffRequest(proto.Message):
     r"""Describes an API diff request.
+
     Attributes:
         threat_type (google.cloud.webrisk_v1beta1.types.ThreatType):
             The ThreatList to update.
@@ -71,6 +72,7 @@ class ComputeThreatListDiffRequest(proto.Message):
 
     class Constraints(proto.Message):
         r"""The constraints for this diff.
+
         Attributes:
             max_diff_entries (int):
                 The maximum size in number of entries. The diff will not
@@ -100,6 +102,7 @@ class Constraints(proto.Message):
 
 class ComputeThreatListDiffResponse(proto.Message):
     r"""
+
     Attributes:
         response_type (google.cloud.webrisk_v1beta1.types.ComputeThreatListDiffResponse.ResponseType):
             The type of response. This may indicate that
@@ -137,6 +140,7 @@ class ResponseType(proto.Enum):
 
     class Checksum(proto.Message):
         r"""The expected state of a client's local database.
+
         Attributes:
             sha256 (bytes):
                 The SHA256 hash of the client state; that is,
@@ -158,6 +162,7 @@ class Checksum(proto.Message):
 
 class SearchUrisRequest(proto.Message):
     r"""Request to check URI entries against threatLists.
+
     Attributes:
         uri (str):
             Required. The URI to be checked for matches.
@@ -171,6 +176,7 @@ class SearchUrisRequest(proto.Message):
 
 class SearchUrisResponse(proto.Message):
     r"""
+
     Attributes:
         threat (google.cloud.webrisk_v1beta1.types.SearchUrisResponse.ThreatUri):
             The threat list matches. This may be empty if
@@ -179,6 +185,7 @@ class SearchUrisResponse(proto.Message):
 
     class ThreatUri(proto.Message):
         r"""Contains threat information on a matching uri.
+
         Attributes:
             threat_types (Sequence[google.cloud.webrisk_v1beta1.types.ThreatType]):
                 The ThreatList this threat belongs to.
@@ -215,6 +222,7 @@ class SearchHashesRequest(proto.Message):
 
 class SearchHashesResponse(proto.Message):
     r"""
+
     Attributes:
         threats (Sequence[google.cloud.webrisk_v1beta1.types.SearchHashesResponse.ThreatHash]):
             The full hashes that matched the requested
@@ -227,6 +235,7 @@ class SearchHashesResponse(proto.Message):
 
     class ThreatHash(proto.Message):
         r"""Contains threat information on a matching hash.
+
         Attributes:
             threat_types (Sequence[google.cloud.webrisk_v1beta1.types.ThreatType]):
                 The ThreatList this threat belongs to.
@@ -276,6 +285,7 @@ class ThreatEntryAdditions(proto.Message):
 
 class ThreatEntryRemovals(proto.Message):
     r"""Contains the set of entries to remove from a local database.
+
     Attributes:
         raw_indices (google.cloud.webrisk_v1beta1.types.RawIndices):
             The raw removal indices for a local list.
@@ -293,6 +303,7 @@ class ThreatEntryRemovals(proto.Message):
 
 class RawIndices(proto.Message):
     r"""A set of raw indices to remove from a local list.
+
     Attributes:
         indices (Sequence[int]):
             The indices to remove from a