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 #129

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 @@ -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(
Expand Down
18 changes: 14 additions & 4 deletions google/cloud/webrisk_v1/services/web_risk_service/client.py
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
Expand Up @@ -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,
Expand Down
Expand Up @@ -356,5 +356,8 @@ def create_submission(
)
return self._stubs["create_submission"]

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


__all__ = ("WebRiskServiceGrpcTransport",)
Expand Up @@ -362,5 +362,8 @@ def create_submission(
)
return self._stubs["create_submission"]

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


__all__ = ("WebRiskServiceGrpcAsyncIOTransport",)
13 changes: 13 additions & 0 deletions google/cloud/webrisk_v1/types/webrisk.py
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
Expand Up @@ -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(
Expand Down
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
Expand Up @@ -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,
Expand Down
Expand Up @@ -312,5 +312,8 @@ def search_hashes(
)
return self._stubs["search_hashes"]

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


__all__ = ("WebRiskServiceV1Beta1GrpcTransport",)
Expand Up @@ -318,5 +318,8 @@ def search_hashes(
)
return self._stubs["search_hashes"]

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


__all__ = ("WebRiskServiceV1Beta1GrpcAsyncIOTransport",)
11 changes: 11 additions & 0 deletions google/cloud/webrisk_v1beta1/types/webrisk.py
Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand Down