feat: Add new gcloud commands, API clients, and third-party libraries across various services.

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/notebooks/__init__.py

@@ -0,0 +1 @@
+

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/notebooks/environments.py

@@ -0,0 +1,129 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""'notebooks environments create' command."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.api_lib.notebooks import util
+from googlecloudsdk.core import log
+from googlecloudsdk.core import resources
+
+
+def CreateEnvironment(args, messages):
+  """Creates the Environment message for the create request.
+
+  Args:
+    args: Argparse object from Command.Run
+    messages: Module containing messages definition for the specified API.
+
+  Returns:
+    Instance of the Environment message.
+  """
+
+  def CreateContainerImageFromArgs(args):
+    container_image = messages.ContainerImage(
+        repository=args.container_repository, tag=args.container_tag)
+    return container_image
+
+  def CreateVmImageFromArgs(args):
+    vm_image = messages.VmImage(
+        project=args.vm_image_project)
+    if args.IsSpecified('vm_image_family'):
+      vm_image.imageFamily = args.vm_image_family
+    else:
+      vm_image.imageName = args.vm_image_name
+    return vm_image
+
+  if args.IsSpecified('vm_image_project'):
+    vm_image = CreateVmImageFromArgs(args)
+  else:
+    container_image = CreateContainerImageFromArgs(args)
+
+  environment = messages.Environment(
+      name=args.environment,
+      description=args.description,
+      displayName=args.display_name,
+      postStartupScript=args.post_startup_script)
+  if args.IsSpecified('vm_image_project'):
+    environment.vmImage = vm_image
+  else:
+    environment.containerImage = container_image
+  return environment
+
+
+def CreateEnvironmentCreateRequest(args, messages):
+  parent = util.GetParentForEnvironment(args)
+  environment = CreateEnvironment(args, messages)
+  return messages.NotebooksProjectsLocationsEnvironmentsCreateRequest(
+      parent=parent, environment=environment, environmentId=args.environment)
+
+
+def CreateEnvironmentListRequest(args, messages):
+  parent = util.GetParentFromArgs(args)
+  return messages.NotebooksProjectsLocationsEnvironmentsListRequest(
+      parent=parent)
+
+
+def CreateEnvironmentDeleteRequest(args, messages):
+  environment = GetEnvironmentResource(args).RelativeName()
+  return messages.NotebooksProjectsLocationsEnvironmentsDeleteRequest(
+      name=environment)
+
+
+def CreateEnvironmentDescribeRequest(args, messages):
+  environment = GetEnvironmentResource(args).RelativeName()
+  return messages.NotebooksProjectsLocationsEnvironmentsGetRequest(
+      name=environment)
+
+
+def GetEnvironmentResource(args):
+  return args.CONCEPTS.environment.Parse()
+
+
+def GetEnvironmentURI(resource):
+  environment = resources.REGISTRY.ParseRelativeName(
+      resource.name, collection='notebooks.projects.locations.environments')
+  return environment.SelfLink()
+
+
+def HandleLRO(operation,
+              args,
+              environment_service,
+              release_track,
+              is_delete=False):
+  """Handles Long Running Operations for both cases of async."""
+  logging_method = log.DeletedResource if is_delete else log.CreatedResource
+  if args.async_:
+    logging_method(
+        util.GetOperationResource(operation.name, release_track),
+        kind='notebooks environment {0}'.format(args.environment),
+        is_async=True)
+    return operation
+  else:
+    response = util.WaitForOperation(
+        operation,
+        'Waiting for Environment [{}] to be {} with [{}]'.format(
+            args.environment, 'deleted' if is_delete else 'created',
+            operation.name),
+        service=environment_service,
+        release_track=release_track,
+        is_delete=is_delete)
+    logging_method(
+        util.GetOperationResource(operation.name, release_track),
+        kind='notebooks environment {0}'.format(args.environment),
+        is_async=False)
+    return response

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/notebooks/instances.py

@@ -0,0 +1,470 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""notebooks instances api helper."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import enum
+from googlecloudsdk.api_lib.notebooks import environments as env_util
+from googlecloudsdk.api_lib.notebooks import util
+from googlecloudsdk.command_lib.util.apis import arg_utils
+from googlecloudsdk.core import log
+from googlecloudsdk.core import resources
+
+_RESERVATION_AFFINITY_KEY = 'compute.googleapis.com/reservation-name'
+
+
+def CreateInstance(args, client, messages):
+  """Creates the Instance message for the create request.
+
+  Args:
+    args: Argparse object from Command.Run
+    client(base_api.BaseApiClient): An instance of the specified API client.
+    messages: Module containing messages definition for the specified API.
+
+  Returns:
+    Instance of the Instance message.
+  """
+
+  def GetContainerImageFromExistingEnvironment():
+    environment_service = client.projects_locations_environments
+    result = environment_service.Get(
+        env_util.CreateEnvironmentDescribeRequest(args, messages))
+    return result.containerImage
+
+  def GetVmImageFromExistingEnvironment():
+    environment_service = client.projects_locations_environments
+    result = environment_service.Get(
+        env_util.CreateEnvironmentDescribeRequest(args, messages))
+    return result.vmImage
+
+  def GetKmsRelativeName():
+    if args.IsSpecified('kms_key'):
+      return args.CONCEPTS.kms_key.Parse().RelativeName()
+
+  def GetNetworkRelativeName():
+    if args.IsSpecified('network'):
+      return args.CONCEPTS.network.Parse().RelativeName()
+
+  def GetSubnetRelativeName():
+    if args.IsSpecified('subnet'):
+      return args.CONCEPTS.subnet.Parse().RelativeName()
+
+  def CreateAcceleratorConfigMessage():
+    accelerator_config = messages.AcceleratorConfig
+    type_enum = None
+    if args.IsSpecified('accelerator_type'):
+      type_enum = arg_utils.ChoiceEnumMapper(
+          arg_name='accelerator-type',
+          message_enum=accelerator_config.TypeValueValuesEnum,
+          include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+              arg_utils.EnumNameToChoice(args.accelerator_type))
+    return accelerator_config(
+        type=type_enum, coreCount=args.accelerator_core_count)
+
+  def GetBootDisk():
+    type_enum = None
+    if args.IsSpecified('boot_disk_type'):
+      instance_message = messages.Instance
+      type_enum = arg_utils.ChoiceEnumMapper(
+          arg_name='boot-disk-type',
+          message_enum=instance_message.BootDiskTypeValueValuesEnum,
+          include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+              arg_utils.EnumNameToChoice(args.boot_disk_type))
+    return type_enum
+
+  def GetDataDisk():
+    type_enum = None
+    if args.IsSpecified('data_disk_type'):
+      instance_message = messages.Instance
+      type_enum = arg_utils.ChoiceEnumMapper(
+          arg_name='data-disk-type',
+          message_enum=instance_message.DataDiskTypeValueValuesEnum,
+          include_filter=lambda x: 'UNSPECIFIED' not in x,
+      ).GetEnumForChoice(arg_utils.EnumNameToChoice(args.data_disk_type))
+    return type_enum
+
+  def GetDiskEncryption():
+    type_enum = None
+    if args.IsSpecified('disk_encryption'):
+      instance_message = messages.Instance
+      type_enum = arg_utils.ChoiceEnumMapper(
+          arg_name='disk-encryption',
+          message_enum=instance_message.DiskEncryptionValueValuesEnum,
+          include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+              arg_utils.EnumNameToChoice(args.disk_encryption))
+    return type_enum
+
+  def CreateContainerImageFromArgs():
+    if args.IsSpecified('environment'):
+      return GetContainerImageFromExistingEnvironment()
+    if args.IsSpecified('container_repository'):
+      container_image = messages.ContainerImage(
+          repository=args.container_repository, tag=args.container_tag)
+      return container_image
+    return None
+
+  def CreateVmImageFromArgs():
+    """Create VmImage Message from an environment or from args."""
+    if args.IsSpecified('environment'):
+      return GetVmImageFromExistingEnvironment()
+    if args.IsSpecified('container_repository'):
+      return None
+    vm_image = messages.VmImage(project=args.vm_image_project)
+    if args.IsSpecified('vm_image_name'):
+      vm_image.imageName = args.vm_image_name
+    else:
+      vm_image.imageFamily = args.vm_image_family
+    return vm_image
+
+  def GetInstanceOwnersFromArgs():
+    if args.IsSpecified('instance_owners'):
+      return [args.instance_owners]
+    return []
+
+  def GetLabelsFromArgs():
+    if args.IsSpecified('labels'):
+      labels_message = messages.Instance.LabelsValue
+      return labels_message(additionalProperties=[
+          labels_message.AdditionalProperty(key=key, value=value)
+          for key, value in args.labels.items()
+      ])
+    return None
+
+  def GetMetadataFromArgs():
+    if args.IsSpecified('metadata'):
+      metadata_message = messages.Instance.MetadataValue
+      return metadata_message(additionalProperties=[
+          metadata_message.AdditionalProperty(key=key, value=value)
+          for key, value in args.metadata.items()
+      ])
+    return None
+
+  def GetShieldedInstanceConfigFromArgs():
+    if not (args.IsSpecified('shielded_vm_secure_boot') or
+            args.IsSpecified('shielded_vm_vtpm') or
+            args.IsSpecified('shielded_vm_integrity_monitoring')):
+      return None
+    shielded_instance_config_message = messages.ShieldedInstanceConfig
+    return shielded_instance_config_message(
+        enableIntegrityMonitoring=args.shielded_vm_integrity_monitoring,
+        enableSecureBoot=args.shielded_vm_secure_boot,
+        enableVtpm=args.shielded_vm_vtpm,
+    )
+
+  def GetTagsFromArgs():
+    if args.IsSpecified('tags'):
+      return args.tags
+    return []
+
+  # TODO(b/194714138): Consider adding validation for specific reservation.
+  def GetReservationAffinityConfigFromArgs():
+    if not (args.IsSpecified('reservation_affinity') or
+            args.IsSpecified('reservation')):
+      return None
+
+    def GetReservationAffinityEnum():
+      type_enum = None
+      if args.IsSpecified('reservation_affinity'):
+        reservation_affinity_message = messages.ReservationAffinity
+        type_enum = arg_utils.ChoiceEnumMapper(
+            arg_name='reservation-affinity',
+            message_enum=(reservation_affinity_message
+                          .ConsumeReservationTypeValueValuesEnum),
+            include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+                arg_utils.EnumNameToChoice(args.reservation_affinity))
+      return type_enum
+
+    reservation_affinity_enum = GetReservationAffinityEnum()
+    reservation_key = None
+    reservation_values = []
+
+    if (reservation_affinity_enum == messages.ReservationAffinity
+        .ConsumeReservationTypeValueValuesEnum.SPECIFIC_RESERVATION):
+      # Currently, the key is fixed and the value is the name of the
+      # reservation.
+      # The value being a repeated field is reserved for future use when user
+      # can specify more than one reservation names from which the VM can take
+      # capacity from.
+      reservation_key = _RESERVATION_AFFINITY_KEY
+      reservation_values = [args.reservation]
+
+    reservation_config_message = messages.ReservationAffinity
+    return reservation_config_message(
+        consumeReservationType=reservation_affinity_enum,
+        key=reservation_key,
+        values=reservation_values,
+    )
+
+  instance = messages.Instance(
+      name=args.instance,
+      postStartupScript=args.post_startup_script,
+      customGpuDriverPath=args.custom_gpu_driver_path,
+      instanceOwners=GetInstanceOwnersFromArgs(),
+      kmsKey=GetKmsRelativeName(),
+      machineType=args.machine_type,
+      network=GetNetworkRelativeName(),
+      noProxyAccess=args.no_proxy_access,
+      noPublicIp=args.no_public_ip,
+      serviceAccount=args.service_account,
+      subnet=GetSubnetRelativeName(),
+      vmImage=CreateVmImageFromArgs(),
+      acceleratorConfig=CreateAcceleratorConfigMessage(),
+      bootDiskType=GetBootDisk(),
+      bootDiskSizeGb=args.boot_disk_size,
+      dataDiskType=GetDataDisk(),
+      dataDiskSizeGb=args.data_disk_size,
+      noRemoveDataDisk=args.no_remove_data_disk,
+      containerImage=CreateContainerImageFromArgs(),
+      diskEncryption=GetDiskEncryption(),
+      labels=GetLabelsFromArgs(),
+      metadata=GetMetadataFromArgs(),
+      installGpuDriver=args.install_gpu_driver,
+      shieldedInstanceConfig=GetShieldedInstanceConfigFromArgs(),
+      reservationAffinity=GetReservationAffinityConfigFromArgs(),
+      tags=GetTagsFromArgs(),
+  )
+  return instance
+
+
+def CreateInstanceCreateRequest(args, client, messages):
+  parent = util.GetParentForInstance(args)
+  instance = CreateInstance(args, client, messages)
+  return messages.NotebooksProjectsLocationsInstancesCreateRequest(
+      parent=parent, instance=instance, instanceId=args.instance)
+
+
+def CreateInstanceListRequest(args, messages):
+  parent = util.GetParentFromArgs(args)
+  return messages.NotebooksProjectsLocationsInstancesListRequest(parent=parent)
+
+
+def CreateInstanceDeleteRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  return messages.NotebooksProjectsLocationsInstancesDeleteRequest(
+      name=instance)
+
+
+def CreateInstanceDescribeRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  return messages.NotebooksProjectsLocationsInstancesGetRequest(name=instance)
+
+
+def CreateInstanceRegisterRequest(args, messages):
+  instance = GetInstanceResource(args)
+  parent = util.GetLocationResource(instance.locationsId,
+                                    instance.projectsId).RelativeName()
+  register_request = messages.RegisterInstanceRequest(
+      instanceId=instance.Name())
+  return messages.NotebooksProjectsLocationsInstancesRegisterRequest(
+      parent=parent, registerInstanceRequest=register_request)
+
+
+def CreateInstanceResetRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  reset_request = messages.ResetInstanceRequest()
+  return messages.NotebooksProjectsLocationsInstancesResetRequest(
+      name=instance, resetInstanceRequest=reset_request)
+
+
+def CreateInstanceStartRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  start_request = messages.StartInstanceRequest()
+  return messages.NotebooksProjectsLocationsInstancesStartRequest(
+      name=instance, startInstanceRequest=start_request)
+
+
+def CreateInstanceStopRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  stop_request = messages.StopInstanceRequest()
+  return messages.NotebooksProjectsLocationsInstancesStopRequest(
+      name=instance, stopInstanceRequest=stop_request)
+
+
+def CreateSetAcceleratorRequest(args, messages):
+  """Create and return Accelerator update request."""
+  instance = GetInstanceResource(args).RelativeName()
+  set_acc_request = messages.SetInstanceAcceleratorRequest()
+  accelerator_config = messages.SetInstanceAcceleratorRequest
+  if args.IsSpecified('accelerator_core_count'):
+    set_acc_request.coreCount = args.accelerator_core_count
+  if args.IsSpecified('accelerator_type'):
+    type_enum = arg_utils.ChoiceEnumMapper(
+        arg_name='accelerator-type',
+        message_enum=accelerator_config.TypeValueValuesEnum,
+        include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+            arg_utils.EnumNameToChoice(args.accelerator_type))
+    set_acc_request.type = type_enum
+  return messages.NotebooksProjectsLocationsInstancesSetAcceleratorRequest(
+      name=instance, setInstanceAcceleratorRequest=set_acc_request)
+
+
+def CreateSetLabelsRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  set_label_request = messages.SetInstanceLabelsRequest()
+  labels_message = messages.SetInstanceLabelsRequest.LabelsValue
+  set_label_request.labels = labels_message(additionalProperties=[
+      labels_message.AdditionalProperty(key=key, value=value)
+      for key, value in args.labels.items()
+  ])
+  return messages.NotebooksProjectsLocationsInstancesSetLabelsRequest(
+      name=instance, setInstanceLabelsRequest=set_label_request)
+
+
+def CreateSetMachineTypeRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  set_machine_request = messages.SetInstanceMachineTypeRequest(
+      machineType=args.machine_type)
+  return messages.NotebooksProjectsLocationsInstancesSetMachineTypeRequest(
+      name=instance, setInstanceMachineTypeRequest=set_machine_request)
+
+
+def CreateInstanceGetHealthRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  return messages.NotebooksProjectsLocationsInstancesGetInstanceHealthRequest(
+      name=instance)
+
+
+def CreateInstanceIsUpgradeableRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  return messages.NotebooksProjectsLocationsInstancesIsUpgradeableRequest(
+      notebookInstance=instance)
+
+
+def CreateInstanceUpgradeRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  upgrade_request = messages.UpgradeInstanceRequest()
+  return messages.NotebooksProjectsLocationsInstancesUpgradeRequest(
+      name=instance, upgradeInstanceRequest=upgrade_request)
+
+
+def CreateInstanceRollbackRequest(args, messages):
+  instance = GetInstanceResource(args).RelativeName()
+  rollback_request = messages.RollbackInstanceRequest(
+      targetSnapshot=args.target_snapshot)
+  return messages.NotebooksProjectsLocationsInstancesRollbackRequest(
+      name=instance, rollbackInstanceRequest=rollback_request)
+
+
+def CreateInstanceDiagnoseRequest(args, messages):
+  """"Create and return Diagnose request."""
+  instance = GetInstanceResource(args).RelativeName()
+  diagnostic_config = messages.DiagnosticConfig(
+      gcsBucket=args.gcs_bucket,
+  )
+  if args.IsSpecified('relative_path'):
+    diagnostic_config.relativePath = args.relative_path
+  if args.IsSpecified('enable-repair'):
+    diagnostic_config.repairFlagEnabled = True
+  if args.IsSpecified('enable-packet-capture'):
+    diagnostic_config.packetCaptureFlagEnabled = True
+  if args.IsSpecified('enable-copy-home-files'):
+    diagnostic_config.copyHomeFilesFlagEnabled = True
+
+  timeout_minutes = None
+  if args.IsSpecified('timeout_minutes'):
+    timeout_minutes = int(args.timeout_minutes)
+
+  return messages.NotebooksProjectsLocationsInstancesDiagnoseRequest(
+      name=instance, diagnoseInstanceRequest=messages.DiagnoseInstanceRequest(
+          diagnosticConfig=diagnostic_config, timeoutMinutes=timeout_minutes))
+
+
+def CreateInstanceMigrateRequest(args, messages):
+  """"Create and return Migrate request."""
+  instance = GetInstanceResource(args).RelativeName()
+
+  def GetPostStartupScriptOption():
+    type_enum = None
+    if args.IsSpecified('post_startup_script_option'):
+      request_message = messages.MigrateInstanceRequest
+      type_enum = arg_utils.ChoiceEnumMapper(
+          arg_name='post-startup-script-option',
+          message_enum=request_message.PostStartupScriptOptionValueValuesEnum,
+          include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+              arg_utils.EnumNameToChoice(args.post_startup_script_option))
+    return type_enum
+
+  return messages.NotebooksProjectsLocationsInstancesMigrateRequest(
+      name=instance,
+      migrateInstanceRequest=messages.MigrateInstanceRequest(
+          postStartupScriptOption=GetPostStartupScriptOption(),
+      ))
+
+
+def GetInstanceResource(args):
+  return args.CONCEPTS.instance.Parse()
+
+
+def GetInstanceURI(resource):
+  instance = resources.REGISTRY.ParseRelativeName(
+      resource.name, collection='notebooks.projects.locations.instances')
+  return instance.SelfLink()
+
+
+class OperationType(enum.Enum):
+  CREATE = (log.CreatedResource, 'created')
+  UPDATE = (log.UpdatedResource, 'updated')
+  UPGRADE = (log.UpdatedResource, 'upgraded')
+  ROLLBACK = (log.UpdatedResource, 'rolled back')
+  DELETE = (log.DeletedResource, 'deleted')
+  RESET = (log.ResetResource, 'reset')
+  MIGRATE = (log.UpdatedResource, 'migrated')
+
+
+def HandleLRO(operation,
+              args,
+              instance_service,
+              release_track,
+              operation_type=OperationType.UPDATE):
+  """Handles Long Running Operations for both cases of async.
+
+  Args:
+    operation: The operation to poll.
+    args: ArgParse instance containing user entered arguments.
+    instance_service: The service to get the resource after the long running
+      operation completes.
+    release_track: base.ReleaseTrack object.
+    operation_type: Enum value of type OperationType indicating the kind of
+      operation to wait for.
+
+  Raises:
+    apitools.base.py.HttpError: if the request returns an HTTP error
+
+  Returns:
+    The Instance resource if synchronous, else the Operation Resource.
+  """
+  logging_method = operation_type.value[0]
+  if args.async_:
+    logging_method(
+        util.GetOperationResource(operation.name, release_track),
+        kind='notebooks instance {0}'.format(args.instance),
+        is_async=True)
+    return operation
+  else:
+    response = util.WaitForOperation(
+        operation,
+        'Waiting for operation on Instance [{}] to be {} with [{}]'.format(
+            args.instance, operation_type.value[1], operation.name),
+        service=instance_service,
+        release_track=release_track,
+        is_delete=(operation_type.value[1] == 'deleted'))
+    logging_method(
+        util.GetOperationResource(operation.name, release_track),
+        kind='notebooks instance {0}'.format(args.instance),
+        is_async=False)
+    return response

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/notebooks/locations.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""'notebooks environments create' command."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.api_lib.notebooks import util
+
+
+def CreateLocationListRequest(args, messages):
+  project_name = util.GetProjectResource(args.project).RelativeName()
+  return messages.NotebooksProjectsLocationsListRequest(
+      name=project_name)
+
+
+def GetLocationURI(resource):
+  location_resource = util.GetLocationResource(resource.name)
+  return location_resource.SelfLink()

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/notebooks/runtimes.py

@@ -0,0 +1,281 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2022 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""notebooks runtimes api helper."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+import enum
+from googlecloudsdk.api_lib.notebooks import util
+from googlecloudsdk.command_lib.util.apis import arg_utils
+from googlecloudsdk.core import log
+from googlecloudsdk.core import resources
+
+
+def CreateRuntime(args, messages):
+  """Creates the Runtime message for the create request.
+
+  Args:
+    args: Argparse object from Command.Run
+    messages: Module containing messages definition for the specified API.
+
+  Returns:
+    Runtime of the Runtime message.
+  """
+
+  def GetRuntimeVirtualMachineFromArgs():
+    machine_type = 'n1-standard-4'
+    if args.IsSpecified('machine_type'):
+      machine_type = args.machine_type
+    virtual_machine_config = messages.VirtualMachineConfig(
+        machineType=machine_type, dataDisk=messages.LocalDisk())
+    return messages.VirtualMachine(virtualMachineConfig=virtual_machine_config)
+
+  def GetRuntimeAccessConfigFromArgs():
+    runtime_access_config = messages.RuntimeAccessConfig
+    type_enum = None
+    if args.IsSpecified('runtime_access_type'):
+      type_enum = arg_utils.ChoiceEnumMapper(
+          arg_name='runtime-access-type',
+          message_enum=runtime_access_config.AccessTypeValueValuesEnum,
+          include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+              arg_utils.EnumNameToChoice(args.runtime_access_type))
+    return runtime_access_config(
+        accessType=type_enum, runtimeOwner=args.runtime_owner)
+
+  def GetPostStartupScriptBehavior():
+    type_enum = None
+    if args.IsSpecified('post_startup_script_behavior'):
+      runtime_software_config_message = messages.RuntimeSoftwareConfig
+      type_enum = arg_utils.ChoiceEnumMapper(
+          arg_name='post-startup-script-behavior',
+          message_enum=(runtime_software_config_message
+                        .PostStartupScriptBehaviorTypeValueValuesEnum),
+          include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+              arg_utils.EnumNameToChoice(args.post_startup_script_behavior))
+    return type_enum
+
+  def GetRuntimeSoftwareConfigFromArgs():
+    runtime_software_config = messages.RuntimeSoftwareConfig()
+    if args.IsSpecified('idle_shutdown_timeout'):
+      runtime_software_config.idleShutdownTimeout = args.idle_shutdown_timeout
+    if args.IsSpecified('install_gpu_driver'):
+      runtime_software_config.installGpuDriver = args.install_gpu_driver
+    if args.IsSpecified('custom_gpu_driver_path'):
+      runtime_software_config.customGpuDriverPath = args.custom_gpu_driver_path
+    if args.IsSpecified('post_startup_script'):
+      runtime_software_config.postStartupScript = args.post_startup_script
+    if args.IsSpecified('post_startup_script_behavior'):
+      runtime_software_config.postStartupScriptBehavior = (
+          GetPostStartupScriptBehavior())
+    return runtime_software_config
+
+  runtime = messages.Runtime(
+      name=args.runtime,
+      virtualMachine=GetRuntimeVirtualMachineFromArgs(),
+      accessConfig=GetRuntimeAccessConfigFromArgs(),
+      softwareConfig=GetRuntimeSoftwareConfigFromArgs(),
+  )
+  return runtime
+
+
+def CreateRuntimeCreateRequest(args, messages):
+  parent = util.GetParentForRuntime(args)
+  runtime = CreateRuntime(args, messages)
+  return messages.NotebooksProjectsLocationsRuntimesCreateRequest(
+      parent=parent, runtime=runtime, runtimeId=args.runtime)
+
+
+def CreateRuntimeListRequest(args, messages):
+  parent = util.GetParentFromArgs(args)
+  return messages.NotebooksProjectsLocationsRuntimesListRequest(parent=parent)
+
+
+def GetRuntimeResource(args):
+  return args.CONCEPTS.runtime.Parse()
+
+
+def GetSwitchRuntimeRequest(args, messages):
+  """Create and return switch runtime request."""
+  machine_type = 'n1-standard-4'
+  if args.IsSpecified('machine_type'):
+    machine_type = args.machine_type
+  runtime_accelerator_config = messages.RuntimeAcceleratorConfig()
+  if args.IsSpecified('accelerator_core_count'):
+    runtime_accelerator_config.coreCount = args.accelerator_core_count
+  if args.IsSpecified('accelerator_type'):
+    type_enum = arg_utils.ChoiceEnumMapper(
+        arg_name='accelerator-type',
+        message_enum=runtime_accelerator_config.TypeValueValuesEnum,
+        include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+            arg_utils.EnumNameToChoice(args.accelerator_type))
+    runtime_accelerator_config.type = type_enum
+  return messages.SwitchRuntimeRequest(
+      machineType=machine_type, acceleratorConfig=runtime_accelerator_config)
+
+
+def CreateRuntimeDeleteRequest(args, messages):
+  runtime = GetRuntimeResource(args).RelativeName()
+  return messages.NotebooksProjectsLocationsRuntimesDeleteRequest(
+      name=runtime)
+
+
+def CreateRuntimeResetRequest(args, messages):
+  runtime = GetRuntimeResource(args).RelativeName()
+  reset_request = messages.ResetRuntimeRequest()
+  return messages.NotebooksProjectsLocationsRuntimesResetRequest(
+      name=runtime, resetRuntimeRequest=reset_request)
+
+
+def CreateRuntimeStartRequest(args, messages):
+  runtime = GetRuntimeResource(args).RelativeName()
+  start_request = messages.StartRuntimeRequest()
+  return messages.NotebooksProjectsLocationsRuntimesStartRequest(
+      name=runtime, startRuntimeRequest=start_request)
+
+
+def CreateRuntimeStopRequest(args, messages):
+  runtime = GetRuntimeResource(args).RelativeName()
+  stop_request = messages.StopRuntimeRequest()
+  return messages.NotebooksProjectsLocationsRuntimesStopRequest(
+      name=runtime, stopRuntimeRequest=stop_request)
+
+
+def CreateRuntimeSwitchRequest(args, messages):
+  runtime = GetRuntimeResource(args).RelativeName()
+  switch_request = GetSwitchRuntimeRequest(args, messages)
+  return messages.NotebooksProjectsLocationsRuntimesSwitchRequest(
+      name=runtime, switchRuntimeRequest=switch_request)
+
+
+def CreateRuntimeDescribeRequest(args, messages):
+  runtime = GetRuntimeResource(args).RelativeName()
+  return messages.NotebooksProjectsLocationsRuntimesGetRequest(name=runtime)
+
+
+def CreateRuntimeDiagnoseRequest(args, messages):
+  """"Create and return Diagnose request."""
+  runtime = GetRuntimeResource(args).RelativeName()
+  diagnostic_config = messages.DiagnosticConfig(
+      gcsBucket=args.gcs_bucket,
+  )
+  if args.IsSpecified('relative_path'):
+    diagnostic_config.relativePath = args.relative_path
+  if args.IsSpecified('enable-repair'):
+    diagnostic_config.repairFlagEnabled = True
+  if args.IsSpecified('enable-packet-capture'):
+    diagnostic_config.packetCaptureFlagEnabled = True
+  if args.IsSpecified('enable-copy-home-files'):
+    diagnostic_config.copyHomeFilesFlagEnabled = True
+
+  timeout_minutes = None
+  if args.IsSpecified('timeout_minutes'):
+    timeout_minutes = int(args.timeout_minutes)
+
+  return messages.NotebooksProjectsLocationsRuntimesDiagnoseRequest(
+      name=runtime, diagnoseRuntimeRequest=messages.DiagnoseRuntimeRequest(
+          diagnosticConfig=diagnostic_config, timeoutMinutes=timeout_minutes))
+
+
+def CreateRuntimeMigrateRequest(args, messages):
+  """"Create and return Migrate request."""
+  runtime = GetRuntimeResource(args).RelativeName()
+
+  def GetNetworkRelativeName():
+    if args.IsSpecified('network'):
+      return args.CONCEPTS.network.Parse().RelativeName()
+
+  def GetSubnetRelativeName():
+    if args.IsSpecified('subnet'):
+      return args.CONCEPTS.subnet.Parse().RelativeName()
+
+  def GetPostStartupScriptOption():
+    type_enum = None
+    if args.IsSpecified('post_startup_script_option'):
+      request_message = messages.MigrateRuntimeRequest
+      type_enum = arg_utils.ChoiceEnumMapper(
+          arg_name='post-startup-script-option',
+          message_enum=request_message.PostStartupScriptOptionValueValuesEnum,
+          include_filter=lambda x: 'UNSPECIFIED' not in x).GetEnumForChoice(
+              arg_utils.EnumNameToChoice(args.post_startup_script_option))
+    return type_enum
+
+  return messages.NotebooksProjectsLocationsRuntimesMigrateRequest(
+      name=runtime,
+      migrateRuntimeRequest=messages.MigrateRuntimeRequest(
+          network=GetNetworkRelativeName(),
+          subnet=GetSubnetRelativeName(),
+          serviceAccount=args.service_account,
+          postStartupScriptOption=GetPostStartupScriptOption(),
+      ))
+
+
+def GetRuntimeURI(resource):
+  instance = resources.REGISTRY.ParseRelativeName(
+      resource.name, collection='notebooks.projects.locations.runtimes')
+  return instance.SelfLink()
+
+
+class OperationType(enum.Enum):
+  CREATE = (log.CreatedResource, 'created')
+  DELETE = (log.DeletedResource, 'deleted')
+  UPDATE = (log.UpdatedResource, 'updated')
+  RESET = (log.ResetResource, 'reset')
+  MIGRATE = (log.UpdatedResource, 'migrated')
+
+
+def HandleLRO(operation,
+              args,
+              runtime_service,
+              release_track,
+              operation_type=OperationType.CREATE):
+  """Handles Long Running Operations for both cases of async.
+
+  Args:
+    operation: The operation to poll.
+    args: ArgParse instance containing user entered arguments.
+    runtime_service: The service to get the resource after the long running
+      operation completes.
+    release_track: base.ReleaseTrack object.
+    operation_type: Enum value of type OperationType indicating the kind of
+      operation to wait for.
+
+  Raises:
+    apitools.base.py.HttpError: if the request returns an HTTP error
+
+  Returns:
+    The Runtime resource if synchronous, else the Operation Resource.
+  """
+  logging_method = operation_type.value[0]
+  if args.async_:
+    logging_method(
+        util.GetOperationResource(operation.name, release_track),
+        kind='notebooks runtime {0}'.format(args.runtime),
+        is_async=True)
+    return operation
+  else:
+    response = util.WaitForOperation(
+        operation,
+        'Waiting for operation on Runtime [{}] to be {} with [{}]'.format(
+            args.runtime, operation_type.value[1], operation.name),
+        service=runtime_service,
+        release_track=release_track,
+        is_delete=(operation_type.value[1] == 'deleted'))
+    logging_method(
+        util.GetOperationResource(operation.name, release_track),
+        kind='notebooks runtime {0}'.format(args.runtime),
+        is_async=False)
+    return response

gcloud auth application-default login/google-cloud-sdk/lib/googlecloudsdk/api_lib/notebooks/util.py

@@ -0,0 +1,134 @@
+# -*- coding: utf-8 -*- #
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A library that is used to support trace commands."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import unicode_literals
+
+from googlecloudsdk.api_lib.util import apis as core_apis
+from googlecloudsdk.api_lib.util import waiter
+from googlecloudsdk.calliope import base
+from googlecloudsdk.core import properties
+from googlecloudsdk.core import resources
+
+
+def ApiVersionSelector(release_track):
+  """Returns the correct API version.
+
+  Args:
+    release_track: base.ReleaseTrack object
+  """
+  if release_track == base.ReleaseTrack.GA:
+    return 'v1'
+  return 'v1beta1'
+
+
+def GetClient(release_track=base.ReleaseTrack.GA):
+  """Returns the client for the trace API."""
+  return core_apis.GetClientInstance('notebooks',
+                                     ApiVersionSelector(release_track))
+
+
+def GetMessages(release_track=base.ReleaseTrack.GA):
+  """Returns the messages for the trace API."""
+  return core_apis.GetMessagesModule('notebooks',
+                                     ApiVersionSelector(release_track))
+
+
+def GetLocationResource(location, project=None):
+  if not project:
+    project = properties.VALUES.core.project.Get(required=True)
+  return resources.REGISTRY.Parse(
+      location,
+      params={'projectsId': project},
+      collection='notebooks.projects.locations')
+
+
+def GetParentForRuntime(args):
+  if args.IsSpecified('runtime'):
+    runtime = args.CONCEPTS.runtime.Parse()
+    return GetLocationResource(runtime.locationsId,
+                               runtime.projectsId).RelativeName()
+
+
+def GetParentForInstance(args):
+  if args.IsSpecified('instance'):
+    instance = args.CONCEPTS.instance.Parse()
+    return GetLocationResource(instance.locationsId,
+                               instance.projectsId).RelativeName()
+
+
+def GetParentForEnvironment(args):
+  if args.IsSpecified('environment'):
+    environment = args.CONCEPTS.environment.Parse()
+    return GetLocationResource(environment.locationsId,
+                               environment.projectsId).RelativeName()
+
+
+def GetProjectResource(project):
+  """Returns the resource for the current project."""
+  return resources.REGISTRY.Parse(
+      project or properties.VALUES.core.project.Get(required=True),
+      collection='notebooks.projects')
+
+
+def GetParentFromArgs(args):
+  project = properties.VALUES.core.project.Get(required=True)
+  if args.IsSpecified('location'):
+    return GetLocationResource(args.location, project).RelativeName()
+  elif properties.VALUES.notebooks.location.IsExplicitlySet():
+    return GetLocationResource(
+        properties.VALUES.notebooks.location.Get(required=True),
+        project).RelativeName()
+
+
+def GetOperationResource(name, release_track=None):
+  return resources.REGISTRY.ParseRelativeName(
+      name,
+      collection='notebooks.projects.locations.operations',
+      api_version=ApiVersionSelector(release_track))
+
+
+def WaitForOperation(operation,
+                     message,
+                     service,
+                     release_track,
+                     is_delete=False):
+  """Waits for the given google.longrunning.Operation to complete.
+
+  Args:
+    operation: The operation to poll.
+    message: String to display for default progress_tracker.
+    service: The service to get the resource after the long running operation
+      completes.
+    release_track: base.ReleaseTrack object.
+    is_delete: Bool indicating is Poller should fetch resource post operation.
+
+  Raises:
+    apitools.base.py.HttpError: if the request returns an HTTP error
+
+  Returns:
+    The created Environment resource.
+  """
+  operation_ref = GetOperationResource(operation.name, release_track)
+  client = GetClient(release_track)
+  if is_delete:
+    poller = waiter.CloudOperationPollerNoResources(
+        client.projects_locations_operations)
+  else:
+    poller = waiter.CloudOperationPoller(
+        service, client.projects_locations_operations)
+  return waiter.WaitFor(poller, operation_ref, message)