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

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 @@ -2284,6 +2284,12 @@ async def update_security_marks(
# 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/securitycenter_v1/services/security_center/client.py
Expand Up @@ -455,10 +455,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 create_source(
Expand Down Expand Up @@ -2467,6 +2464,19 @@ def update_security_marks(
# 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 @@ -385,6 +385,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 @@ -920,5 +920,8 @@ def update_security_marks(
)
return self._stubs["update_security_marks"]

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


__all__ = ("SecurityCenterGrpcTransport",)
Expand Up @@ -937,5 +937,8 @@ def update_security_marks(
)
return self._stubs["update_security_marks"]

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


__all__ = ("SecurityCenterGrpcAsyncIOTransport",)
Expand Up @@ -26,6 +26,7 @@

class NotificationMessage(proto.Message):
r"""Cloud SCC's Notification

Attributes:
notification_config_name (str):
Name of the notification config that
Expand Down
Expand Up @@ -43,6 +43,7 @@ class OrganizationSettings(proto.Message):

class AssetDiscoveryConfig(proto.Message):
r"""The configuration used for Asset Discovery runs.

Attributes:
project_ids (Sequence[str]):
The project ids to use for filtering asset
Expand Down
1 change: 1 addition & 0 deletions google/cloud/securitycenter_v1/types/resource.py
Expand Up @@ -25,6 +25,7 @@

class Resource(proto.Message):
r"""Information related to the Google Cloud resource.

Attributes:
name (str):
The full resource name of the resource. See:
Expand Down
Expand Up @@ -25,6 +25,7 @@

class RunAssetDiscoveryResponse(proto.Message):
r"""Response of asset discovery run

Attributes:
state (google.cloud.securitycenter_v1.types.RunAssetDiscoveryResponse.State):
The state of an asset discovery run.
Expand Down
27 changes: 27 additions & 0 deletions google/cloud/securitycenter_v1/types/securitycenter_service.py
Expand Up @@ -68,6 +68,7 @@

class CreateFindingRequest(proto.Message):
r"""Request message for creating a finding.

Attributes:
parent (str):
Required. Resource name of the new finding's parent. Its
Expand All @@ -92,6 +93,7 @@ class CreateFindingRequest(proto.Message):

class CreateNotificationConfigRequest(proto.Message):
r"""Request message for creating a notification config.

Attributes:
parent (str):
Required. Resource name of the new notification config's
Expand All @@ -118,6 +120,7 @@ class CreateNotificationConfigRequest(proto.Message):

class CreateSourceRequest(proto.Message):
r"""Request message for creating a source.

Attributes:
parent (str):
Required. Resource name of the new source's parent. Its
Expand All @@ -134,6 +137,7 @@ class CreateSourceRequest(proto.Message):

class DeleteNotificationConfigRequest(proto.Message):
r"""Request message for deleting a notification config.

Attributes:
name (str):
Required. Name of the notification config to delete. Its
Expand All @@ -146,6 +150,7 @@ class DeleteNotificationConfigRequest(proto.Message):

class GetNotificationConfigRequest(proto.Message):
r"""Request message for getting a notification config.

Attributes:
name (str):
Required. Name of the notification config to get. Its format
Expand All @@ -158,6 +163,7 @@ class GetNotificationConfigRequest(proto.Message):

class GetOrganizationSettingsRequest(proto.Message):
r"""Request message for getting organization settings.

Attributes:
name (str):
Required. Name of the organization to get organization
Expand All @@ -170,6 +176,7 @@ class GetOrganizationSettingsRequest(proto.Message):

class GetSourceRequest(proto.Message):
r"""Request message for getting a source.

Attributes:
name (str):
Required. Relative resource name of the source. Its format
Expand All @@ -181,6 +188,7 @@ class GetSourceRequest(proto.Message):

class GroupAssetsRequest(proto.Message):
r"""Request message for grouping by assets.

Attributes:
parent (str):
Required. Name of the organization to groupBy. Its format is
Expand Down Expand Up @@ -351,6 +359,7 @@ class GroupAssetsRequest(proto.Message):

class GroupAssetsResponse(proto.Message):
r"""Response message for grouping by assets.

Attributes:
group_by_results (Sequence[google.cloud.securitycenter_v1.types.GroupResult]):
Group results. There exists an element for
Expand Down Expand Up @@ -382,6 +391,7 @@ def raw_page(self):

class GroupFindingsRequest(proto.Message):
r"""Request message for grouping by findings.

Attributes:
parent (str):
Required. Name of the source to groupBy. Its format is
Expand Down Expand Up @@ -553,6 +563,7 @@ class GroupFindingsRequest(proto.Message):

class GroupFindingsResponse(proto.Message):
r"""Response message for group by findings.

Attributes:
group_by_results (Sequence[google.cloud.securitycenter_v1.types.GroupResult]):
Group results. There exists an element for
Expand Down Expand Up @@ -603,6 +614,7 @@ class GroupResult(proto.Message):

class ListNotificationConfigsRequest(proto.Message):
r"""Request message for listing notification configs.

Attributes:
parent (str):
Required. Name of the organization to list notification
Expand All @@ -625,6 +637,7 @@ class ListNotificationConfigsRequest(proto.Message):

class ListNotificationConfigsResponse(proto.Message):
r"""Response message for listing notification configs.

Attributes:
notification_configs (Sequence[google.cloud.securitycenter_v1.types.NotificationConfig]):
Notification configs belonging to the
Expand All @@ -646,6 +659,7 @@ def raw_page(self):

class ListSourcesRequest(proto.Message):
r"""Request message for listing sources.

Attributes:
parent (str):
Required. Resource name of the parent of sources to list.
Expand All @@ -669,6 +683,7 @@ class ListSourcesRequest(proto.Message):

class ListSourcesResponse(proto.Message):
r"""Response message for listing sources.

Attributes:
sources (Sequence[google.cloud.securitycenter_v1.types.Source]):
Sources belonging to the requested parent.
Expand All @@ -687,6 +702,7 @@ def raw_page(self):

class ListAssetsRequest(proto.Message):
r"""Request message for listing assets.

Attributes:
parent (str):
Required. Name of the organization assets should belong to.
Expand Down Expand Up @@ -859,6 +875,7 @@ class ListAssetsRequest(proto.Message):

class ListAssetsResponse(proto.Message):
r"""Response message for listing assets.

Attributes:
list_assets_results (Sequence[google.cloud.securitycenter_v1.types.ListAssetsResponse.ListAssetsResult]):
Assets matching the list request.
Expand All @@ -874,6 +891,7 @@ class ListAssetsResponse(proto.Message):

class ListAssetsResult(proto.Message):
r"""Result containing the Asset and its State.

Attributes:
asset (google.cloud.securitycenter_v1.types.Asset):
Asset matching the search request.
Expand Down Expand Up @@ -916,6 +934,7 @@ def raw_page(self):

class ListFindingsRequest(proto.Message):
r"""Request message for listing findings.

Attributes:
parent (str):
Required. Name of the source the findings belong to. Its
Expand Down Expand Up @@ -1088,6 +1107,7 @@ class ListFindingsRequest(proto.Message):

class ListFindingsResponse(proto.Message):
r"""Response message for listing findings.

Attributes:
list_findings_results (Sequence[google.cloud.securitycenter_v1.types.ListFindingsResponse.ListFindingsResult]):
Findings matching the list request.
Expand All @@ -1103,6 +1123,7 @@ class ListFindingsResponse(proto.Message):

class ListFindingsResult(proto.Message):
r"""Result containing the Finding and its StateChange.

Attributes:
finding (google.cloud.securitycenter_v1.types.Finding):
Finding matching the search request.
Expand Down Expand Up @@ -1190,6 +1211,7 @@ def raw_page(self):

class SetFindingStateRequest(proto.Message):
r"""Request message for updating a finding's state.

Attributes:
name (str):
Required. The relative resource name of the finding. See:
Expand Down Expand Up @@ -1223,6 +1245,7 @@ class RunAssetDiscoveryRequest(proto.Message):

class UpdateFindingRequest(proto.Message):
r"""Request message for updating or creating a finding.

Attributes:
finding (google.cloud.securitycenter_v1.types.Finding):
Required. The finding resource to update or create if it
Expand Down Expand Up @@ -1250,6 +1273,7 @@ class UpdateFindingRequest(proto.Message):

class UpdateNotificationConfigRequest(proto.Message):
r"""Request message for updating a notification config.

Attributes:
notification_config (google.cloud.securitycenter_v1.types.NotificationConfig):
Required. The notification config to update.
Expand All @@ -1269,6 +1293,7 @@ class UpdateNotificationConfigRequest(proto.Message):

class UpdateOrganizationSettingsRequest(proto.Message):
r"""Request message for updating an organization's settings.

Attributes:
organization_settings (google.cloud.securitycenter_v1.types.OrganizationSettings):
Required. The organization settings resource
Expand All @@ -1290,6 +1315,7 @@ class UpdateOrganizationSettingsRequest(proto.Message):

class UpdateSourceRequest(proto.Message):
r"""Request message for updating a source.

Attributes:
source (google.cloud.securitycenter_v1.types.Source):
Required. The source resource to update.
Expand All @@ -1307,6 +1333,7 @@ class UpdateSourceRequest(proto.Message):

class UpdateSecurityMarksRequest(proto.Message):
r"""Request message for updating a SecurityMarks resource.

Attributes:
security_marks (google.cloud.securitycenter_v1.types.SecurityMarks):
Required. The security marks resource to
Expand Down
Expand Up @@ -1838,6 +1838,12 @@ async def update_security_marks(
# 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 @@ -423,10 +423,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 create_source(
Expand Down Expand Up @@ -1996,6 +1993,19 @@ def update_security_marks(
# 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 @@ -337,6 +337,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 @@ -772,5 +772,8 @@ def update_security_marks(
)
return self._stubs["update_security_marks"]

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


__all__ = ("SecurityCenterGrpcTransport",)