From 22efa38a5e9be1e1137a68329fa947e8a116753c Mon Sep 17 00:00:00 2001
From: "gcf-owl-bot[bot]" <78513119+gcf-owl-bot[bot]@users.noreply.github.com>
Date: Thu, 7 Oct 2021 00:36:23 +0000
Subject: [PATCH] feat: add context manager support in client (#99)

- [ ] Regenerate this pull request now.

chore: fix docstring for first attribute of protos

committer: @busunkim96
PiperOrigin-RevId: 401271153

Source-Link: https://github.com/googleapis/googleapis/commit/787f8c9a731f44e74a90b9847d48659ca9462d10

Source-Link: https://github.com/googleapis/googleapis-gen/commit/81decffe9fc72396a8153e756d1d67a6eecfd620
Copy-Tag: eyJwIjoiLmdpdGh1Yi8uT3dsQm90LnlhbWwiLCJoIjoiODFkZWNmZmU5ZmM3MjM5NmE4MTUzZTc1NmQxZDY3YTZlZWNmZDYyMCJ9
---
 .../managed_notebook_service/async_client.py  |  6 +++
 .../managed_notebook_service/client.py        | 18 +++++--
 .../transports/base.py                        |  9 ++++
 .../transports/grpc.py                        |  3 ++
 .../transports/grpc_asyncio.py                |  3 ++
 .../services/notebook_service/async_client.py |  6 +++
 .../services/notebook_service/client.py       | 18 +++++--
 .../notebook_service/transports/base.py       |  9 ++++
 .../notebook_service/transports/grpc.py       |  3 ++
 .../transports/grpc_asyncio.py                |  3 ++
 google/cloud/notebooks_v1/types/execution.py  |  3 ++
 .../notebooks_v1/types/instance_config.py     |  1 +
 .../notebooks_v1/types/managed_service.py     | 10 ++++
 google/cloud/notebooks_v1/types/runtime.py    |  5 ++
 google/cloud/notebooks_v1/types/schedule.py   |  1 +
 google/cloud/notebooks_v1/types/service.py    | 37 ++++++++++++++
 .../services/notebook_service/async_client.py |  6 +++
 .../services/notebook_service/client.py       | 18 +++++--
 .../notebook_service/transports/base.py       |  9 ++++
 .../notebook_service/transports/grpc.py       |  3 ++
 .../transports/grpc_asyncio.py                |  3 ++
 .../cloud/notebooks_v1beta1/types/service.py  | 22 ++++++++
 .../test_managed_notebook_service.py          | 50 +++++++++++++++++++
 .../notebooks_v1/test_notebook_service.py     | 50 +++++++++++++++++++
 .../test_notebook_service.py                  | 50 +++++++++++++++++++
 25 files changed, 334 insertions(+), 12 deletions(-)

diff --git a/google/cloud/notebooks_v1/services/managed_notebook_service/async_client.py b/google/cloud/notebooks_v1/services/managed_notebook_service/async_client.py
index 942b7e6..80ca851 100644
--- a/google/cloud/notebooks_v1/services/managed_notebook_service/async_client.py
+++ b/google/cloud/notebooks_v1/services/managed_notebook_service/async_client.py
@@ -938,6 +938,12 @@ async def report_runtime_event(
         # 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/notebooks_v1/services/managed_notebook_service/client.py b/google/cloud/notebooks_v1/services/managed_notebook_service/client.py
index 8750a11..8384122 100644
--- a/google/cloud/notebooks_v1/services/managed_notebook_service/client.py
+++ b/google/cloud/notebooks_v1/services/managed_notebook_service/client.py
@@ -353,10 +353,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_runtimes(
@@ -1126,6 +1123,19 @@ def report_runtime_event(
         # 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/notebooks_v1/services/managed_notebook_service/transports/base.py b/google/cloud/notebooks_v1/services/managed_notebook_service/transports/base.py
index fd4998a..60f8d2f 100644
--- a/google/cloud/notebooks_v1/services/managed_notebook_service/transports/base.py
+++ b/google/cloud/notebooks_v1/services/managed_notebook_service/transports/base.py
@@ -186,6 +186,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."""
diff --git a/google/cloud/notebooks_v1/services/managed_notebook_service/transports/grpc.py b/google/cloud/notebooks_v1/services/managed_notebook_service/transports/grpc.py
index 0ffd08f..a9f6ffc 100644
--- a/google/cloud/notebooks_v1/services/managed_notebook_service/transports/grpc.py
+++ b/google/cloud/notebooks_v1/services/managed_notebook_service/transports/grpc.py
@@ -495,5 +495,8 @@ def report_runtime_event(
             )
         return self._stubs["report_runtime_event"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("ManagedNotebookServiceGrpcTransport",)
diff --git a/google/cloud/notebooks_v1/services/managed_notebook_service/transports/grpc_asyncio.py b/google/cloud/notebooks_v1/services/managed_notebook_service/transports/grpc_asyncio.py
index bd6828a..3aa853e 100644
--- a/google/cloud/notebooks_v1/services/managed_notebook_service/transports/grpc_asyncio.py
+++ b/google/cloud/notebooks_v1/services/managed_notebook_service/transports/grpc_asyncio.py
@@ -513,5 +513,8 @@ def report_runtime_event(
             )
         return self._stubs["report_runtime_event"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("ManagedNotebookServiceGrpcAsyncIOTransport",)
diff --git a/google/cloud/notebooks_v1/services/notebook_service/async_client.py b/google/cloud/notebooks_v1/services/notebook_service/async_client.py
index 3eb38cc..2b7950e 100644
--- a/google/cloud/notebooks_v1/services/notebook_service/async_client.py
+++ b/google/cloud/notebooks_v1/services/notebook_service/async_client.py
@@ -2530,6 +2530,12 @@ async def create_execution(
         # 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/notebooks_v1/services/notebook_service/client.py b/google/cloud/notebooks_v1/services/notebook_service/client.py
index 17d1021..40fb5eb 100644
--- a/google/cloud/notebooks_v1/services/notebook_service/client.py
+++ b/google/cloud/notebooks_v1/services/notebook_service/client.py
@@ -400,10 +400,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_instances(
@@ -2778,6 +2775,19 @@ def create_execution(
         # 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/notebooks_v1/services/notebook_service/transports/base.py b/google/cloud/notebooks_v1/services/notebook_service/transports/base.py
index f248456..2c0a19e 100644
--- a/google/cloud/notebooks_v1/services/notebook_service/transports/base.py
+++ b/google/cloud/notebooks_v1/services/notebook_service/transports/base.py
@@ -270,6 +270,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."""
diff --git a/google/cloud/notebooks_v1/services/notebook_service/transports/grpc.py b/google/cloud/notebooks_v1/services/notebook_service/transports/grpc.py
index d8fe6cc..efb8a86 100644
--- a/google/cloud/notebooks_v1/services/notebook_service/transports/grpc.py
+++ b/google/cloud/notebooks_v1/services/notebook_service/transports/grpc.py
@@ -1099,5 +1099,8 @@ def create_execution(
             )
         return self._stubs["create_execution"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("NotebookServiceGrpcTransport",)
diff --git a/google/cloud/notebooks_v1/services/notebook_service/transports/grpc_asyncio.py b/google/cloud/notebooks_v1/services/notebook_service/transports/grpc_asyncio.py
index 897e685..3eba203 100644
--- a/google/cloud/notebooks_v1/services/notebook_service/transports/grpc_asyncio.py
+++ b/google/cloud/notebooks_v1/services/notebook_service/transports/grpc_asyncio.py
@@ -1146,5 +1146,8 @@ def create_execution(
             )
         return self._stubs["create_execution"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("NotebookServiceGrpcAsyncIOTransport",)
diff --git a/google/cloud/notebooks_v1/types/execution.py b/google/cloud/notebooks_v1/types/execution.py
index 07768b9..56d8eb9 100644
--- a/google/cloud/notebooks_v1/types/execution.py
+++ b/google/cloud/notebooks_v1/types/execution.py
@@ -25,6 +25,7 @@
 
 class ExecutionTemplate(proto.Message):
     r"""The description a notebook execution workload.
+
     Attributes:
         scale_tier (google.cloud.notebooks_v1.types.ExecutionTemplate.ScaleTier):
             Required. Scale tier of the hardware used for
@@ -174,6 +175,7 @@ class SchedulerAcceleratorConfig(proto.Message):
 
     class DataprocParameters(proto.Message):
         r"""Parameters used in Dataproc JobType executions.
+
         Attributes:
             cluster (str):
                 URI for cluster used to run Dataproc execution. Format:
@@ -202,6 +204,7 @@ class DataprocParameters(proto.Message):
 
 class Execution(proto.Message):
     r"""The definition of a single executed notebook.
+
     Attributes:
         execution_template (google.cloud.notebooks_v1.types.ExecutionTemplate):
             execute metadata including name, hardware
diff --git a/google/cloud/notebooks_v1/types/instance_config.py b/google/cloud/notebooks_v1/types/instance_config.py
index e76798a..b89e721 100644
--- a/google/cloud/notebooks_v1/types/instance_config.py
+++ b/google/cloud/notebooks_v1/types/instance_config.py
@@ -23,6 +23,7 @@
 
 class InstanceConfig(proto.Message):
     r"""Notebook instance configurations that can be updated.
+
     Attributes:
         notebook_upgrade_schedule (str):
             Cron expression in UTC timezone, used to schedule instance
diff --git a/google/cloud/notebooks_v1/types/managed_service.py b/google/cloud/notebooks_v1/types/managed_service.py
index 432a7bf..97c5951 100644
--- a/google/cloud/notebooks_v1/types/managed_service.py
+++ b/google/cloud/notebooks_v1/types/managed_service.py
@@ -38,6 +38,7 @@
 
 class ListRuntimesRequest(proto.Message):
     r"""Request for listing Managed Notebook Runtimes.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -56,6 +57,7 @@ class ListRuntimesRequest(proto.Message):
 
 class ListRuntimesResponse(proto.Message):
     r"""Response for listing Managed Notebook Runtimes.
+
     Attributes:
         runtimes (Sequence[google.cloud.notebooks_v1.types.Runtime]):
             A list of returned Runtimes.
@@ -82,6 +84,7 @@ def raw_page(self):
 
 class GetRuntimeRequest(proto.Message):
     r"""Request for getting a Managed Notebook Runtime.
+
     Attributes:
         name (str):
             Required. Format:
@@ -93,6 +96,7 @@ class GetRuntimeRequest(proto.Message):
 
 class CreateRuntimeRequest(proto.Message):
     r"""Request for creating a Managed Notebook Runtime.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -111,6 +115,7 @@ class CreateRuntimeRequest(proto.Message):
 
 class DeleteRuntimeRequest(proto.Message):
     r"""Request for deleting a Managed Notebook Runtime.
+
     Attributes:
         name (str):
             Required. Format:
@@ -122,6 +127,7 @@ class DeleteRuntimeRequest(proto.Message):
 
 class StartRuntimeRequest(proto.Message):
     r"""Request for starting a Managed Notebook Runtime.
+
     Attributes:
         name (str):
             Required. Format:
@@ -133,6 +139,7 @@ class StartRuntimeRequest(proto.Message):
 
 class StopRuntimeRequest(proto.Message):
     r"""Request for stopping a Managed Notebook Runtime.
+
     Attributes:
         name (str):
             Required. Format:
@@ -144,6 +151,7 @@ class StopRuntimeRequest(proto.Message):
 
 class SwitchRuntimeRequest(proto.Message):
     r"""Request for switching a Managed Notebook Runtime.
+
     Attributes:
         name (str):
             Required. Format:
@@ -163,6 +171,7 @@ class SwitchRuntimeRequest(proto.Message):
 
 class ResetRuntimeRequest(proto.Message):
     r"""Request for reseting a Managed Notebook Runtime.
+
     Attributes:
         name (str):
             Required. Format:
@@ -174,6 +183,7 @@ class ResetRuntimeRequest(proto.Message):
 
 class ReportRuntimeEventRequest(proto.Message):
     r"""Request for reporting a Managed Notebook Event.
+
     Attributes:
         name (str):
             Required. Format:
diff --git a/google/cloud/notebooks_v1/types/runtime.py b/google/cloud/notebooks_v1/types/runtime.py
index 64fda82..74104e1 100644
--- a/google/cloud/notebooks_v1/types/runtime.py
+++ b/google/cloud/notebooks_v1/types/runtime.py
@@ -39,6 +39,7 @@
 
 class Runtime(proto.Message):
     r"""The definition of a Runtime for a managed notebook instance.
+
     Attributes:
         name (str):
             Output only. The resource name of the runtime. Format:
@@ -158,6 +159,7 @@ class EncryptionConfig(proto.Message):
 
 class LocalDisk(proto.Message):
     r"""An Local attached disk resource.
+
     Attributes:
         auto_delete (bool):
             Optional. Output only. Specifies whether the
@@ -331,6 +333,7 @@ class DiskType(proto.Enum):
 
 class RuntimeAccessConfig(proto.Message):
     r"""Specifies the login configuration for Runtime
+
     Attributes:
         access_type (google.cloud.notebooks_v1.types.RuntimeAccessConfig.RuntimeAccessType):
             The type of access mode this instance.
@@ -446,6 +449,7 @@ class RuntimeShieldedInstanceConfig(proto.Message):
 
 class VirtualMachine(proto.Message):
     r"""Runtime using Virtual Machine for computing.
+
     Attributes:
         instance_name (str):
             Output only. The user-friendly name of the
@@ -466,6 +470,7 @@ class VirtualMachine(proto.Message):
 
 class VirtualMachineConfig(proto.Message):
     r"""The config settings for virtual machine.
+
     Attributes:
         zone (str):
             Output only. The zone where the virtual machine is located.
diff --git a/google/cloud/notebooks_v1/types/schedule.py b/google/cloud/notebooks_v1/types/schedule.py
index 118f0b9..6722be3 100644
--- a/google/cloud/notebooks_v1/types/schedule.py
+++ b/google/cloud/notebooks_v1/types/schedule.py
@@ -26,6 +26,7 @@
 
 class Schedule(proto.Message):
     r"""The definition of a schedule.
+
     Attributes:
         name (str):
             Output only. The name of this schedule. Format:
diff --git a/google/cloud/notebooks_v1/types/service.py b/google/cloud/notebooks_v1/types/service.py
index c93e08a..d49ad11 100644
--- a/google/cloud/notebooks_v1/types/service.py
+++ b/google/cloud/notebooks_v1/types/service.py
@@ -71,6 +71,7 @@
 
 class OperationMetadata(proto.Message):
     r"""Represents the metadata of the long-running operation.
+
     Attributes:
         create_time (google.protobuf.timestamp_pb2.Timestamp):
             The time the operation was created.
@@ -108,6 +109,7 @@ class OperationMetadata(proto.Message):
 
 class ListInstancesRequest(proto.Message):
     r"""Request for listing notebook instances.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -126,6 +128,7 @@ class ListInstancesRequest(proto.Message):
 
 class ListInstancesResponse(proto.Message):
     r"""Response for listing notebook instances.
+
     Attributes:
         instances (Sequence[google.cloud.notebooks_v1.types.Instance]):
             A list of returned instances.
@@ -152,6 +155,7 @@ def raw_page(self):
 
 class GetInstanceRequest(proto.Message):
     r"""Request for getting a notebook instance.
+
     Attributes:
         name (str):
             Required. Format:
@@ -163,6 +167,7 @@ class GetInstanceRequest(proto.Message):
 
 class CreateInstanceRequest(proto.Message):
     r"""Request for creating a notebook instance.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -181,6 +186,7 @@ class CreateInstanceRequest(proto.Message):
 
 class RegisterInstanceRequest(proto.Message):
     r"""Request for registering a notebook instance.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -199,6 +205,7 @@ class RegisterInstanceRequest(proto.Message):
 
 class SetInstanceAcceleratorRequest(proto.Message):
     r"""Request for setting instance accelerator.
+
     Attributes:
         name (str):
             Required. Format:
@@ -222,6 +229,7 @@ class SetInstanceAcceleratorRequest(proto.Message):
 
 class SetInstanceMachineTypeRequest(proto.Message):
     r"""Request for setting instance machine type.
+
     Attributes:
         name (str):
             Required. Format:
@@ -237,6 +245,7 @@ class SetInstanceMachineTypeRequest(proto.Message):
 
 class UpdateInstanceConfigRequest(proto.Message):
     r"""Request for updating instance configurations.
+
     Attributes:
         name (str):
             Required. Format:
@@ -253,6 +262,7 @@ class UpdateInstanceConfigRequest(proto.Message):
 
 class SetInstanceLabelsRequest(proto.Message):
     r"""Request for setting instance labels.
+
     Attributes:
         name (str):
             Required. Format:
@@ -288,6 +298,7 @@ class UpdateShieldedInstanceConfigRequest(proto.Message):
 
 class DeleteInstanceRequest(proto.Message):
     r"""Request for deleting a notebook instance.
+
     Attributes:
         name (str):
             Required. Format:
@@ -299,6 +310,7 @@ class DeleteInstanceRequest(proto.Message):
 
 class StartInstanceRequest(proto.Message):
     r"""Request for starting a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -310,6 +322,7 @@ class StartInstanceRequest(proto.Message):
 
 class StopInstanceRequest(proto.Message):
     r"""Request for stopping a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -321,6 +334,7 @@ class StopInstanceRequest(proto.Message):
 
 class ResetInstanceRequest(proto.Message):
     r"""Request for reseting a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -355,6 +369,7 @@ class ReportInstanceInfoRequest(proto.Message):
 
 class IsInstanceUpgradeableRequest(proto.Message):
     r"""Request for checking if a notebook instance is upgradeable.
+
     Attributes:
         notebook_instance (str):
             Required. Format:
@@ -366,6 +381,7 @@ class IsInstanceUpgradeableRequest(proto.Message):
 
 class IsInstanceUpgradeableResponse(proto.Message):
     r"""Response for checking if a notebook instance is upgradeable.
+
     Attributes:
         upgradeable (bool):
             If an instance is upgradeable.
@@ -390,6 +406,7 @@ class IsInstanceUpgradeableResponse(proto.Message):
 
 class GetInstanceHealthRequest(proto.Message):
     r"""Request for checking if a notebook instance is healthy.
+
     Attributes:
         name (str):
             Required. Format:
@@ -401,6 +418,7 @@ class GetInstanceHealthRequest(proto.Message):
 
 class GetInstanceHealthResponse(proto.Message):
     r"""Response for checking if a notebook instance is healthy.
+
     Attributes:
         health_state (google.cloud.notebooks_v1.types.GetInstanceHealthResponse.HealthState):
             Output only. Runtime health_state.
@@ -426,6 +444,7 @@ class HealthState(proto.Enum):
 
 class UpgradeInstanceRequest(proto.Message):
     r"""Request for upgrading a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -437,6 +456,7 @@ class UpgradeInstanceRequest(proto.Message):
 
 class RollbackInstanceRequest(proto.Message):
     r"""Request for rollbacking a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -453,6 +473,7 @@ class RollbackInstanceRequest(proto.Message):
 
 class UpgradeInstanceInternalRequest(proto.Message):
     r"""Request for upgrading a notebook instance from within the VM
+
     Attributes:
         name (str):
             Required. Format:
@@ -470,6 +491,7 @@ class UpgradeInstanceInternalRequest(proto.Message):
 
 class ListEnvironmentsRequest(proto.Message):
     r"""Request for listing environments.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -488,6 +510,7 @@ class ListEnvironmentsRequest(proto.Message):
 
 class ListEnvironmentsResponse(proto.Message):
     r"""Response for listing environments.
+
     Attributes:
         environments (Sequence[google.cloud.notebooks_v1.types.Environment]):
             A list of returned environments.
@@ -512,6 +535,7 @@ def raw_page(self):
 
 class GetEnvironmentRequest(proto.Message):
     r"""Request for getting a notebook environment.
+
     Attributes:
         name (str):
             Required. Format:
@@ -523,6 +547,7 @@ class GetEnvironmentRequest(proto.Message):
 
 class CreateEnvironmentRequest(proto.Message):
     r"""Request for creating a notebook environment.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -546,6 +571,7 @@ class CreateEnvironmentRequest(proto.Message):
 
 class DeleteEnvironmentRequest(proto.Message):
     r"""Request for deleting a notebook environment.
+
     Attributes:
         name (str):
             Required. Format:
@@ -557,6 +583,7 @@ class DeleteEnvironmentRequest(proto.Message):
 
 class ListSchedulesRequest(proto.Message):
     r"""Request for listing scheduled notebook job.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -581,6 +608,7 @@ class ListSchedulesRequest(proto.Message):
 
 class ListSchedulesResponse(proto.Message):
     r"""Response for listing scheduled notebook job.
+
     Attributes:
         schedules (Sequence[google.cloud.notebooks_v1.types.Schedule]):
             A list of returned instances.
@@ -610,6 +638,7 @@ def raw_page(self):
 
 class GetScheduleRequest(proto.Message):
     r"""Request for getting scheduled notebook.
+
     Attributes:
         name (str):
             Required. Format:
@@ -621,6 +650,7 @@ class GetScheduleRequest(proto.Message):
 
 class DeleteScheduleRequest(proto.Message):
     r"""Request for deleting an Schedule
+
     Attributes:
         name (str):
             Required. Format:
@@ -632,6 +662,7 @@ class DeleteScheduleRequest(proto.Message):
 
 class CreateScheduleRequest(proto.Message):
     r"""Request for created scheduled notebooks
+
     Attributes:
         parent (str):
             Required. Format:
@@ -650,6 +681,7 @@ class CreateScheduleRequest(proto.Message):
 
 class TriggerScheduleRequest(proto.Message):
     r"""Request for created scheduled notebooks
+
     Attributes:
         name (str):
             Required. Format:
@@ -661,6 +693,7 @@ class TriggerScheduleRequest(proto.Message):
 
 class ListExecutionsRequest(proto.Message):
     r"""Request for listing scheduled notebook executions.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -687,6 +720,7 @@ class ListExecutionsRequest(proto.Message):
 
 class ListExecutionsResponse(proto.Message):
     r"""Response for listing scheduled notebook executions
+
     Attributes:
         executions (Sequence[google.cloud.notebooks_v1.types.Execution]):
             A list of returned instances.
@@ -716,6 +750,7 @@ def raw_page(self):
 
 class GetExecutionRequest(proto.Message):
     r"""Request for getting scheduled notebook execution
+
     Attributes:
         name (str):
             Required. Format:
@@ -727,6 +762,7 @@ class GetExecutionRequest(proto.Message):
 
 class DeleteExecutionRequest(proto.Message):
     r"""Request for deleting a scheduled notebook execution
+
     Attributes:
         name (str):
             Required. Format:
@@ -738,6 +774,7 @@ class DeleteExecutionRequest(proto.Message):
 
 class CreateExecutionRequest(proto.Message):
     r"""Request to create notebook execution
+
     Attributes:
         parent (str):
             Required. Format:
diff --git a/google/cloud/notebooks_v1beta1/services/notebook_service/async_client.py b/google/cloud/notebooks_v1beta1/services/notebook_service/async_client.py
index 18b6efd..80127fb 100644
--- a/google/cloud/notebooks_v1beta1/services/notebook_service/async_client.py
+++ b/google/cloud/notebooks_v1beta1/services/notebook_service/async_client.py
@@ -1303,6 +1303,12 @@ async def delete_environment(
         # 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/notebooks_v1beta1/services/notebook_service/client.py b/google/cloud/notebooks_v1beta1/services/notebook_service/client.py
index 764f0a8..9550e37 100644
--- a/google/cloud/notebooks_v1beta1/services/notebook_service/client.py
+++ b/google/cloud/notebooks_v1beta1/services/notebook_service/client.py
@@ -362,10 +362,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_instances(
@@ -1525,6 +1522,19 @@ def delete_environment(
         # 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/notebooks_v1beta1/services/notebook_service/transports/base.py b/google/cloud/notebooks_v1beta1/services/notebook_service/transports/base.py
index 29aa9f7..913d62c 100644
--- a/google/cloud/notebooks_v1beta1/services/notebook_service/transports/base.py
+++ b/google/cloud/notebooks_v1beta1/services/notebook_service/transports/base.py
@@ -225,6 +225,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."""
diff --git a/google/cloud/notebooks_v1beta1/services/notebook_service/transports/grpc.py b/google/cloud/notebooks_v1beta1/services/notebook_service/transports/grpc.py
index a7d327c..710767e 100644
--- a/google/cloud/notebooks_v1beta1/services/notebook_service/transports/grpc.py
+++ b/google/cloud/notebooks_v1beta1/services/notebook_service/transports/grpc.py
@@ -750,5 +750,8 @@ def delete_environment(
             )
         return self._stubs["delete_environment"]
 
+    def close(self):
+        self.grpc_channel.close()
+
 
 __all__ = ("NotebookServiceGrpcTransport",)
diff --git a/google/cloud/notebooks_v1beta1/services/notebook_service/transports/grpc_asyncio.py b/google/cloud/notebooks_v1beta1/services/notebook_service/transports/grpc_asyncio.py
index 88014f3..c3956be 100644
--- a/google/cloud/notebooks_v1beta1/services/notebook_service/transports/grpc_asyncio.py
+++ b/google/cloud/notebooks_v1beta1/services/notebook_service/transports/grpc_asyncio.py
@@ -780,5 +780,8 @@ def delete_environment(
             )
         return self._stubs["delete_environment"]
 
+    def close(self):
+        return self.grpc_channel.close()
+
 
 __all__ = ("NotebookServiceGrpcAsyncIOTransport",)
diff --git a/google/cloud/notebooks_v1beta1/types/service.py b/google/cloud/notebooks_v1beta1/types/service.py
index e25802a..4c987b1 100644
--- a/google/cloud/notebooks_v1beta1/types/service.py
+++ b/google/cloud/notebooks_v1beta1/types/service.py
@@ -52,6 +52,7 @@
 
 class OperationMetadata(proto.Message):
     r"""Represents the metadata of the long-running operation.
+
     Attributes:
         create_time (google.protobuf.timestamp_pb2.Timestamp):
             The time the operation was created.
@@ -89,6 +90,7 @@ class OperationMetadata(proto.Message):
 
 class ListInstancesRequest(proto.Message):
     r"""Request for listing notebook instances.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -107,6 +109,7 @@ class ListInstancesRequest(proto.Message):
 
 class ListInstancesResponse(proto.Message):
     r"""Response for listing notebook instances.
+
     Attributes:
         instances (Sequence[google.cloud.notebooks_v1beta1.types.Instance]):
             A list of returned instances.
@@ -133,6 +136,7 @@ def raw_page(self):
 
 class GetInstanceRequest(proto.Message):
     r"""Request for getting a notebook instance.
+
     Attributes:
         name (str):
             Required. Format:
@@ -144,6 +148,7 @@ class GetInstanceRequest(proto.Message):
 
 class CreateInstanceRequest(proto.Message):
     r"""Request for creating a notebook instance.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -162,6 +167,7 @@ class CreateInstanceRequest(proto.Message):
 
 class RegisterInstanceRequest(proto.Message):
     r"""Request for registering a notebook instance.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -180,6 +186,7 @@ class RegisterInstanceRequest(proto.Message):
 
 class SetInstanceAcceleratorRequest(proto.Message):
     r"""Request for setting instance accelerator.
+
     Attributes:
         name (str):
             Required. Format:
@@ -203,6 +210,7 @@ class SetInstanceAcceleratorRequest(proto.Message):
 
 class SetInstanceMachineTypeRequest(proto.Message):
     r"""Request for setting instance machine type.
+
     Attributes:
         name (str):
             Required. Format:
@@ -218,6 +226,7 @@ class SetInstanceMachineTypeRequest(proto.Message):
 
 class SetInstanceLabelsRequest(proto.Message):
     r"""Request for setting instance labels.
+
     Attributes:
         name (str):
             Required. Format:
@@ -234,6 +243,7 @@ class SetInstanceLabelsRequest(proto.Message):
 
 class DeleteInstanceRequest(proto.Message):
     r"""Request for deleting a notebook instance.
+
     Attributes:
         name (str):
             Required. Format:
@@ -245,6 +255,7 @@ class DeleteInstanceRequest(proto.Message):
 
 class StartInstanceRequest(proto.Message):
     r"""Request for starting a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -256,6 +267,7 @@ class StartInstanceRequest(proto.Message):
 
 class StopInstanceRequest(proto.Message):
     r"""Request for stopping a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -267,6 +279,7 @@ class StopInstanceRequest(proto.Message):
 
 class ResetInstanceRequest(proto.Message):
     r"""Request for reseting a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -301,6 +314,7 @@ class ReportInstanceInfoRequest(proto.Message):
 
 class IsInstanceUpgradeableRequest(proto.Message):
     r"""Request for checking if a notebook instance is upgradeable.
+
     Attributes:
         notebook_instance (str):
             Required. Format:
@@ -312,6 +326,7 @@ class IsInstanceUpgradeableRequest(proto.Message):
 
 class IsInstanceUpgradeableResponse(proto.Message):
     r"""Response for checking if a notebook instance is upgradeable.
+
     Attributes:
         upgradeable (bool):
             If an instance is upgradeable.
@@ -330,6 +345,7 @@ class IsInstanceUpgradeableResponse(proto.Message):
 
 class UpgradeInstanceRequest(proto.Message):
     r"""Request for upgrading a notebook instance
+
     Attributes:
         name (str):
             Required. Format:
@@ -341,6 +357,7 @@ class UpgradeInstanceRequest(proto.Message):
 
 class UpgradeInstanceInternalRequest(proto.Message):
     r"""Request for upgrading a notebook instance from within the VM
+
     Attributes:
         name (str):
             Required. Format:
@@ -358,6 +375,7 @@ class UpgradeInstanceInternalRequest(proto.Message):
 
 class ListEnvironmentsRequest(proto.Message):
     r"""Request for listing environments.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -376,6 +394,7 @@ class ListEnvironmentsRequest(proto.Message):
 
 class ListEnvironmentsResponse(proto.Message):
     r"""Response for listing environments.
+
     Attributes:
         environments (Sequence[google.cloud.notebooks_v1beta1.types.Environment]):
             A list of returned environments.
@@ -400,6 +419,7 @@ def raw_page(self):
 
 class GetEnvironmentRequest(proto.Message):
     r"""Request for getting a notebook environment.
+
     Attributes:
         name (str):
             Required. Format:
@@ -411,6 +431,7 @@ class GetEnvironmentRequest(proto.Message):
 
 class CreateEnvironmentRequest(proto.Message):
     r"""Request for creating a notebook environment.
+
     Attributes:
         parent (str):
             Required. Format:
@@ -434,6 +455,7 @@ class CreateEnvironmentRequest(proto.Message):
 
 class DeleteEnvironmentRequest(proto.Message):
     r"""Request for deleting a notebook environment.
+
     Attributes:
         name (str):
             Required. Format:
diff --git a/tests/unit/gapic/notebooks_v1/test_managed_notebook_service.py b/tests/unit/gapic/notebooks_v1/test_managed_notebook_service.py
index 7b89c19..2f53443 100644
--- a/tests/unit/gapic/notebooks_v1/test_managed_notebook_service.py
+++ b/tests/unit/gapic/notebooks_v1/test_managed_notebook_service.py
@@ -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.notebooks_v1.services.managed_notebook_service import (
@@ -2681,6 +2682,9 @@ def test_managed_notebook_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):
@@ -3198,3 +3202,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 = ManagedNotebookServiceAsyncClient(
+        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 = ManagedNotebookServiceClient(
+            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 = ManagedNotebookServiceClient(
+            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()
diff --git a/tests/unit/gapic/notebooks_v1/test_notebook_service.py b/tests/unit/gapic/notebooks_v1/test_notebook_service.py
index 79c2d3a..63e1599 100644
--- a/tests/unit/gapic/notebooks_v1/test_notebook_service.py
+++ b/tests/unit/gapic/notebooks_v1/test_notebook_service.py
@@ -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.notebooks_v1.services.notebook_service import (
@@ -7083,6 +7084,9 @@ def test_notebook_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):
@@ -7664,3 +7668,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 = NotebookServiceAsyncClient(
+        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 = NotebookServiceClient(
+            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 = NotebookServiceClient(
+            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()
diff --git a/tests/unit/gapic/notebooks_v1beta1/test_notebook_service.py b/tests/unit/gapic/notebooks_v1beta1/test_notebook_service.py
index d5e54da..0f45c66 100644
--- a/tests/unit/gapic/notebooks_v1beta1/test_notebook_service.py
+++ b/tests/unit/gapic/notebooks_v1beta1/test_notebook_service.py
@@ -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.notebooks_v1beta1.services.notebook_service import (
@@ -3672,6 +3673,9 @@ def test_notebook_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):
@@ -4205,3 +4209,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 = NotebookServiceAsyncClient(
+        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 = NotebookServiceClient(
+            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 = NotebookServiceClient(
+            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()