def Create(environment_ref, flags):
  """Calls the Composer Environments.Create method.

  Args:
    environment_ref: Resource, the Composer environment resource to create.
    flags: CreateEnvironmentFlags, the flags provided for environment creation.

  Returns:
    Operation: the operation corresponding to the creation of the environment
  """
  messages = api_util.GetMessagesModule(release_track=flags.release_track)
  # Builds environment message and attaches the configuration
  environment = messages.Environment(name=environment_ref.RelativeName())
  environment.config = _CreateConfig(messages, flags)

  if flags.labels:
    environment.labels = api_util.DictToMessage(
        flags.labels, messages.Environment.LabelsValue)

  try:
    return GetService(release_track=flags.release_track).Create(
        api_util.GetMessagesModule(release_track=flags.release_track)
        .ComposerProjectsLocationsEnvironmentsCreateRequest(
            environment=environment,
            parent=environment_ref.Parent().RelativeName()))
  except apitools_exceptions.HttpForbiddenError as e:
    raise exceptions.HttpException(
        e,
        error_format=(
            'Creation operation failed because of lack of proper '
            'permissions. Please, refer to '
            'https://cloud.google.com/composer/docs/how-to/managing/creating '
            'and Composer Creation Troubleshooting pages to resolve this issue.'
        ))
def List(location_refs, page_size, limit=sys.maxsize):
  """Lists Composer Operations across all locations.

  Uses a hardcoded list of locations, as there is way to dynamically
  discover the list of supported locations. Support for new locations
  will be aligned with Cloud SDK releases.

  Args:
    location_refs: [core.resources.Resource], a list of resource reference to
        locations in which to list operations.
    page_size: An integer specifying the maximum number of resources to be
      returned in a single list call.
    limit: An integer specifying the maximum number of operations to list.
        None if all available operations should be returned.

  Returns:
    list: a generator over Operations within the locations in `location_refs`.
  """
  return api_util.AggregateListResults(
      api_util.GetMessagesModule()
      .ComposerProjectsLocationsOperationsListRequest,
      GetService(),
      location_refs,
      'operations',
      page_size,
      limit=limit,
      location_attribute='name')
Exemple #3
0
    def Run(self, args):
        running_state = (api_util.GetMessagesModule(
            release_track=self.ReleaseTrack()).Environment.
                         StateValueValuesEnum.RUNNING)

        env_ref = args.CONCEPTS.environment.Parse()
        env_obj = environments_api_util.Get(env_ref,
                                            release_track=self.ReleaseTrack())

        if env_obj.state != running_state:
            raise command_util.Error(
                'Cannot execute subcommand for environment in state {}. '
                'Must be RUNNING.'.format(env_obj.state))

        cluster_id = env_obj.config.gkeCluster
        cluster_location_id = command_util.ExtractGkeClusterLocationId(env_obj)

        with command_util.TemporaryKubeconfig(cluster_location_id, cluster_id):
            pod = command_util.GetGkePod(pod_substr=WORKER_POD_SUBSTR)

            self.BypassConfirmationPrompt(args)
            kubectl_args = [
                'exec', pod, '-tic', WORKER_CONTAINER, 'airflow',
                args.subcommand
            ]
            if args.cmd_args:
                # Add '--' to the argument list so kubectl won't eat the command args.
                kubectl_args.extend(['--'] + args.cmd_args)
            command_util.RunKubectlCommand(kubectl_args,
                                           out_func=log.status.Print)
def _ConstructAutoscalingPatch(autoscaling_maximum_cpu,
                               autoscaling_minimum_cpu,
                               autoscaling_maximum_memory,
                               autoscaling_minimum_memory, release_track):
    """Constructs an environment patch for Airflow web server machine type.

  Args:
    autoscaling_maximum_cpu: int or None, maximum number of CPU allowed in the
      Environment. Cannot be specified unless autoscaling is enabled.
    autoscaling_minimum_cpu: int or None, minimum number of CPU allowed in the
      Environment. Cannot be specified unless autoscaling is enabled.
    autoscaling_maximum_memory: int or None, maximum memory (in GB) allowed in
      the Environment. Cannot be specified unless autoscaling is enabled.
    autoscaling_minimum_memory: int or None, minimum memory (in GB) allowed in
      the Environment. Cannot be specified unless autoscaling is enabled.
    release_track: base.ReleaseTrack, the release track of command. It dictates
      which Composer client library is used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule(release_track=release_track)
    config = messages.EnvironmentConfig(
        autoscalingConfig=messages.AutoscalingConfig(
            maximumCpu=autoscaling_maximum_cpu,
            maximumMemory=autoscaling_maximum_memory,
            minimumCpu=autoscaling_minimum_cpu,
            minimumMemory=autoscaling_minimum_memory))
    return 'config.autoscaling_config', messages.Environment(config=config)
Exemple #5
0
def _ConstructAirflowConfigsPatch(clear_airflow_configs,
                                  remove_airflow_configs,
                                  update_airflow_configs,
                                  release_track=base.ReleaseTrack.GA):
  """Constructs an environment patch for updating Airflow configs.

  Args:
    clear_airflow_configs: bool, whether to clear the Airflow configs
      dictionary.
    remove_airflow_configs: iterable(string), Iterable of Airflow config
      property names to remove.
    update_airflow_configs: {string: string}, dict of Airflow config property
      names and values to set.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
        which Composer client library will be used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
  messages = api_util.GetMessagesModule(release_track=release_track)
  env_cls = messages.Environment
  airflow_config_overrides_cls = (
      messages.SoftwareConfig.AirflowConfigOverridesValue)
  entry_cls = airflow_config_overrides_cls.AdditionalProperty

  def _BuildEnv(entries):
    software_config = messages.SoftwareConfig(
        airflowConfigOverrides=airflow_config_overrides_cls(
            additionalProperties=entries))
    config = messages.EnvironmentConfig(softwareConfig=software_config)
    return env_cls(config=config)

  return command_util.BuildPartialUpdate(
      clear_airflow_configs, remove_airflow_configs, update_airflow_configs,
      'config.software_config.airflow_config_overrides', entry_cls, _BuildEnv)
Exemple #6
0
def List(location_refs,
         page_size,
         limit=None,
         release_track=base.ReleaseTrack.GA):
    """Lists Composer Environments across all locations.

  Uses a hardcoded list of locations, as there is no way to dynamically
  discover the list of supported locations. Support for new locations
  will be aligned with Cloud SDK releases.

  Args:
    location_refs: [core.resources.Resource], a list of resource reference to
        locations in which to list environments.
    page_size: An integer specifying the maximum number of resources to be
      returned in a single list call.
    limit: An integer specifying the maximum number of environments to list.
        None if all available environments should be returned.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
        which Composer client library will be used.

  Returns:
    list: a generator over Environments in the locations in `location_refs`
  """
    return api_util.AggregateListResults(
        api_util.GetMessagesModule(release_track=release_track).
        ComposerProjectsLocationsEnvironmentsListRequest,
        GetService(release_track=release_track),
        location_refs,
        'environments',
        page_size,
        limit=limit)
def _ConstructPyPiPackagesPatch(clear_pypi_packages,
                                remove_pypi_packages,
                                update_pypi_packages,
                                release_track=base.ReleaseTrack.GA):
    """Constructs an environment patch for partially updating PyPI packages.

  Args:
    clear_pypi_packages: bool, whether to clear the PyPI packages dictionary.
    remove_pypi_packages: iterable(string), Iterable of PyPI package names to
      remove.
    update_pypi_packages: {string: string}, dict mapping PyPI package name
      to optional extras and version specifier.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
        which Composer client library will be used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule(release_track=release_track)
    env_cls = messages.Environment
    pypi_packages_cls = (messages.SoftwareConfig.PypiPackagesValue)
    entry_cls = pypi_packages_cls.AdditionalProperty

    def _BuildEnv(entries):
        software_config = messages.SoftwareConfig(
            pypiPackages=pypi_packages_cls(additionalProperties=entries))
        config = messages.EnvironmentConfig(softwareConfig=software_config)
        return env_cls(config=config)

    return command_util.BuildPartialUpdate(
        clear_pypi_packages, remove_pypi_packages, update_pypi_packages,
        'config.software_config.pypi_packages', entry_cls, _BuildEnv)
def Patch(environment_ref,
          environment_patch,
          update_mask,
          release_track=base.ReleaseTrack.GA):
    """Calls the Composer Environments.Update method.

  Args:
    environment_ref: Resource, the Composer environment resource to update.
    environment_patch: The Environment message specifying the patch associated
      with the update_mask.
    update_mask: A field mask defining the patch.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
      which Composer client library will be used.

  Returns:
    Operation: the operation corresponding to the environment update
  """
    try:
        return GetService(release_track=release_track).Patch(
            api_util.GetMessagesModule(release_track=release_track).
            ComposerProjectsLocationsEnvironmentsPatchRequest(
                name=environment_ref.RelativeName(),
                environment=environment_patch,
                updateMask=update_mask))
    except apitools_exceptions.HttpForbiddenError as e:
        raise exceptions.HttpException(
            e,
            error_format=
            ('Update operation failed because of lack of proper '
             'permissions. Please, refer to '
             'https://cloud.google.com/composer/docs/how-to/managing/updating '
             'and Composer Update Troubleshooting pages to resolve this issue.'
             ))
def _ConstructLabelsPatch(clear_labels,
                          remove_labels,
                          update_labels,
                          release_track=base.ReleaseTrack.GA):
    """Constructs an environment patch for updating labels.

  Args:
    clear_labels: bool, whether to clear the labels dictionary.
    remove_labels: iterable(string), Iterable of label names to remove.
    update_labels: {string: string}, dict of label names and values to set.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
        which Composer client library will be used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule(release_track=release_track)
    env_cls = messages.Environment
    entry_cls = env_cls.LabelsValue.AdditionalProperty

    def _BuildEnv(entries):
        return env_cls(labels=env_cls.LabelsValue(
            additionalProperties=entries))

    return command_util.BuildPartialUpdate(clear_labels, remove_labels,
                                           update_labels, 'labels', entry_cls,
                                           _BuildEnv)
def LoadSnapshot(environment_ref,
                 skip_pypi_packages_installation,
                 snapshot_path,
                 release_track=base.ReleaseTrack.ALPHA):
    """Calls the Composer Environments.LoadSnapshot method.

  Args:
    environment_ref: Resource, the Composer environment resource to Load the
      snapshot for.
    skip_pypi_packages_installation: skip installing the pypi packages during
      the operation.
    snapshot_path: path of the specific snapshot to load the snapshot.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
      which Composer client library will be used.

  Returns:
    Operation: the operation corresponding to loading the snapshot.
  """
    message_module = api_util.GetMessagesModule(release_track=release_track)
    request_message = message_module.ComposerProjectsLocationsEnvironmentsLoadSnapshotRequest(
        environment=environment_ref.RelativeName(),
        loadSnapshotRequest=message_module.LoadSnapshotRequest(
            skipPypiPackagesInstallation=skip_pypi_packages_installation,
            snapshotPath=snapshot_path))
    return GetService(
        release_track=release_track).LoadSnapshot(request_message)
def _ConstructAutoscalingPatch(scheduler_cpu, worker_cpu, scheduler_memory_gb,
                               worker_memory_gb, worker_min_count,
                               worker_max_count, release_track):
    """Constructs an environment patch for Airflow web server machine type.

  Args:
    scheduler_cpu: float or None, CPU allocated to Airflow scheduler.
        Can be specified only in Composer 2.0.0.
    worker_cpu: float or None, CPU allocated to each Airflow worker.
        Can be specified only in Composer 2.0.0.
    scheduler_memory_gb: float or None, memory allocated to Airflow scheduler.
        Can be specified only in Composer 2.0.0.
    worker_memory_gb: float or None, memory allocated to each Airflow worker.
        Can be specified only in Composer 2.0.0.
    worker_min_count: int or None, minimum number of workers
        in the Environment. Can be specified only in Composer 2.0.0.
    worker_max_count: int or None, maximumn number of workers
        in the Environment. Can be specified only in Composer 2.0.0.
    release_track: base.ReleaseTrack, the release track of command. It dictates
      which Composer client library is used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule(release_track=release_track)
    config = messages.EnvironmentConfig(
        workloadsConfig=messages.WorkloadsConfig(
            scheduler=messages.SchedulerResource(cpu=scheduler_cpu,
                                                 memoryGb=scheduler_memory_gb),
            worker=messages.WorkerResource(cpu=worker_cpu,
                                           memoryGb=worker_memory_gb,
                                           minCount=worker_min_count,
                                           maxCount=worker_max_count)))
    mask = ''
    if scheduler_cpu:
        mask = mask + 'config.workloads_config.scheduler.cpu'
    if worker_cpu:
        if mask:
            mask = mask + ','
        mask = mask + 'config.workloads_config.worker.cpu'
    if scheduler_memory_gb:
        if mask:
            mask = mask + ','
        mask = mask + 'config.workloads_config.scheduler.memory_gb'
    if worker_memory_gb:
        if mask:
            mask = mask + ','
        mask = mask + 'config.workloads_config.worker.memory_gb'
    if worker_min_count:
        if mask:
            mask = mask + ','
        mask = mask + 'config.workloads_config.worker.min_count'
    if worker_max_count:
        if mask:
            mask = mask + ','
        mask = mask + 'config.workloads_config.worker.max_count'
    if (scheduler_cpu and worker_cpu and scheduler_memory_gb
            and worker_memory_gb and worker_min_count and worker_max_count):
        mask = 'config.workloads_config'
    return mask, messages.Environment(config=config)
def _ConstructMaintenanceWindowPatch(maintenance_window_start,
                                     maintenance_window_end,
                                     maintenance_window_recurrence,
                                     release_track=base.ReleaseTrack.GA):
    """Constructs an environment patch for updating maintenance window.

  Args:
    maintenance_window_start: Datetime or None, a starting date of the
      maintenance window.
    maintenance_window_end: Datetime or None, an ending date of the maintenance
      window.
    maintenance_window_recurrence: str or None, recurrence RRULE for the
      maintenance window.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
      which Composer client library will be used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule(release_track=release_track)

    window_value = messages.MaintenanceWindow(
        startTime=maintenance_window_start.isoformat(),
        endTime=maintenance_window_end.isoformat(),
        recurrence=maintenance_window_recurrence)
    config = messages.EnvironmentConfig(maintenanceWindow=window_value)

    return 'config.maintenance_window', messages.Environment(config=config)
    def Run(self, args):
        self.DeprecationWarningPrompt(args)
        self.CheckForRequiredCmdArgs(args)

        running_state = (api_util.GetMessagesModule(
            release_track=self.ReleaseTrack()).Environment.
                         StateValueValuesEnum.RUNNING)

        env_ref = args.CONCEPTS.environment.Parse()
        env_obj = environments_api_util.Get(env_ref,
                                            release_track=self.ReleaseTrack())

        if env_obj.state != running_state:
            raise command_util.Error(
                'Cannot execute subcommand for environment in state {}. '
                'Must be RUNNING.'.format(env_obj.state))

        cluster_id = env_obj.config.gkeCluster
        cluster_location_id = command_util.ExtractGkeClusterLocationId(env_obj)

        tty = 'no-tty' not in args

        with command_util.TemporaryKubeconfig(cluster_location_id, cluster_id):
            try:
                image_version = env_obj.config.softwareConfig.imageVersion
                airflow_version = self._ExtractAirflowVersion(image_version)

                self.CheckSubcommandAirflowSupport(args, airflow_version)
                self.CheckSubcommandNestedAirflowSupport(args, airflow_version)

                kubectl_ns = command_util.FetchKubectlNamespace(image_version)
                pod = command_util.GetGkePod(pod_substr=WORKER_POD_SUBSTR,
                                             kubectl_namespace=kubectl_ns)

                log.status.Print(
                    'Executing within the following Kubernetes cluster namespace: '
                    '{}'.format(kubectl_ns))

                self.BypassConfirmationPrompt(args, airflow_version)
                kubectl_args = ['exec', pod, '--stdin']
                if tty:
                    kubectl_args.append('--tty')
                kubectl_args.extend([
                    '--container', WORKER_CONTAINER, '--', 'airflow',
                    args.subcommand
                ])
                if args.subcommand_nested:
                    kubectl_args.append(args.subcommand_nested)
                if args.cmd_args:
                    kubectl_args.extend(args.cmd_args)

                command_util.RunKubectlCommand(
                    command_util.AddKubectlNamespace(kubectl_ns, kubectl_args),
                    out_func=log.out.Print)
            except command_util.KubectlError as e:
                raise self.ConvertKubectlError(e, env_obj)
Exemple #14
0
def IsInRunningState(environment):
  """Returns whether an environment currently is in the RUNNING state.

  Args:
    environment: Environment, an object returned by an API call representing
        the environment to check.
  """
  running_state = (
      api_util.GetMessagesModule().Environment.StateValueValuesEnum.RUNNING)

  return environment.state == running_state
Exemple #15
0
def _ConstructNodeCountPatch(node_count):
    """Constructs an environment patch for node count.

  Args:
    node_count: int, the desired node count

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule()
    config = messages.EnvironmentConfig(nodeCount=node_count)
    return 'config.node_count', messages.Environment(config=config)
def _ConstructEnvVariablesPatch(env_ref,
                                clear_env_variables,
                                remove_env_variables,
                                update_env_variables,
                                release_track=base.ReleaseTrack.GA):
    """Constructs an environment patch for updating environment variables.

  Note that environment variable updates do not support partial update masks
  unlike other map updates due to comments in (b/78298321). For this reason, we
  need to retrieve the Environment, apply an update on EnvVariable dictionary,
  and patch the entire dictionary. The potential race condition here
  (environment variables being updated between when we retrieve them and when we
  send patch request)is not a concern since environment variable updates take
  5 mins to complete, and environments cannot be updated while already in the
  updating state.

  Args:
    env_ref: resource argument, Environment resource argument for environment
      being updated.
    clear_env_variables: bool, whether to clear the environment variables
      dictionary.
    remove_env_variables: iterable(string), Iterable of environment variable
      names to remove.
    update_env_variables: {string: string}, dict of environment variable names
      and values to set.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
        which Composer client library will be used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    env_obj = environments_api_util.Get(env_ref, release_track=release_track)
    initial_env_var_value = env_obj.config.softwareConfig.envVariables
    initial_env_var_list = (initial_env_var_value.additionalProperties
                            if initial_env_var_value else [])

    messages = api_util.GetMessagesModule(release_track=release_track)
    env_cls = messages.Environment
    env_variables_cls = messages.SoftwareConfig.EnvVariablesValue
    entry_cls = env_variables_cls.AdditionalProperty

    def _BuildEnv(entries):
        software_config = messages.SoftwareConfig(
            envVariables=env_variables_cls(additionalProperties=entries))
        config = messages.EnvironmentConfig(softwareConfig=software_config)
        return env_cls(config=config)

    return ('config.software_config.env_variables',
            command_util.BuildFullMapUpdate(clear_env_variables,
                                            remove_env_variables,
                                            update_env_variables,
                                            initial_env_var_list, entry_cls,
                                            _BuildEnv))
def Delete(environment_ref):
    """Calls the Composer Environments.Delete method.

  Args:
    environment_ref: Resource, the Composer environment resource to
        delete.

  Returns:
    Operation: the operation corresponding to the deletion of the environment
  """
    return GetService().Delete(api_util.GetMessagesModule(
    ).ComposerProjectsLocationsEnvironmentsDeleteRequest(
        name=environment_ref.RelativeName()))
Exemple #18
0
def IsInRunningState(environment, release_track=base.ReleaseTrack.GA):
    """Returns whether an environment currently is in the RUNNING state.

  Args:
    environment: Environment, an object returned by an API call representing
        the environment to check.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
        which Composer client library will be used.
  """
    running_state = (api_util.GetMessagesModule(
        release_track=release_track).Environment.StateValueValuesEnum.RUNNING)

    return environment.state == running_state
def Get(environment_ref):
    """Calls the Composer Environments.Get method.

  Args:
    environment_ref: Resource, the Composer environment resource to
        retrieve.

  Returns:
    Environment: the requested environment
  """
    return GetService().Get(api_util.GetMessagesModule().
                            ComposerProjectsLocationsEnvironmentsGetRequest(
                                name=environment_ref.RelativeName()))
def Get(operation_resource):
  """Calls the Composer Operations.Get method.

  Args:
    operation_resource: Resource, the Composer operation resource to
        retrieve.

  Returns:
    Operation: the requested operation
  """
  return GetService().Get(api_util.GetMessagesModule()
                          .ComposerProjectsLocationsOperationsGetRequest(
                              name=operation_resource.RelativeName()))
def Delete(operation_resource):
  """Calls the Composer Operations.Delete method.

  Args:
    operation_resource: Resource, the Composer operation resource to
        delete.

  Returns:
    Empty
  """
  return GetService().Delete(api_util.GetMessagesModule()
                             .ComposerProjectsLocationsOperationsDeleteRequest(
                                 name=operation_resource.RelativeName()))
def _ConstructNodeCountPatch(node_count, release_track=base.ReleaseTrack.GA):
    """Constructs an environment patch for node count.

  Args:
    node_count: int, the desired node count
    release_track: base.ReleaseTrack, the release track of command. Will dictate
        which Composer client library will be used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule(release_track=release_track)
    config = messages.EnvironmentConfig(nodeCount=node_count)
    return 'config.node_count', messages.Environment(config=config)
def _ConstructEnvironmentSizePatch(environment_size,
                                   release_track=base.ReleaseTrack.GA):
    """Constructs an environment patch for environment size.

  Args:
    environment_size: str, the desired environment size.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
      which Composer client library will be used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule(release_track=release_track)
    config = messages.EnvironmentConfig(environmentSize=environment_size)
    return 'config.environment_size', messages.Environment(config=config)
def Delete(environment_ref, release_track=base.ReleaseTrack.GA):
    """Calls the Composer Environments.Delete method.

  Args:
    environment_ref: Resource, the Composer environment resource to delete.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
      which Composer client library will be used.

  Returns:
    Operation: the operation corresponding to the deletion of the environment
  """
    return GetService(release_track=release_track).Delete(
        api_util.GetMessagesModule(release_track=release_track).
        ComposerProjectsLocationsEnvironmentsDeleteRequest(
            name=environment_ref.RelativeName()))
def Get(environment_ref, release_track=base.ReleaseTrack.GA):
    """Calls the Composer Environments.Get method.

  Args:
    environment_ref: Resource, the Composer environment resource to retrieve.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
      which Composer client library will be used.

  Returns:
    Environment: the requested environment
  """
    return GetService(release_track=release_track).Get(
        api_util.GetMessagesModule(release_track=release_track).
        ComposerProjectsLocationsEnvironmentsGetRequest(
            name=environment_ref.RelativeName()))
def Patch(environment_ref, environment_patch, update_mask):
    """Calls the Composer Environments.Update method.

  Args:
    environment_ref: Resource, the Composer environment resource to update.
    environment_patch: The Environment message specifying the patch associated
      with the update_mask.
    update_mask: A field mask defining the patch.
  Returns:
    Operation: the operation corresponding to the environment update
  """
    return GetService().Patch(api_util.GetMessagesModule(
    ).ComposerProjectsLocationsEnvironmentsPatchRequest(
        name=environment_ref.RelativeName(),
        environment=environment_patch,
        updateMask=update_mask))
def Get(operation_resource, release_track=base.ReleaseTrack.GA):
    """Calls the Composer Operations.Get method.

  Args:
    operation_resource: Resource, the Composer operation resource to
        retrieve.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
        which Composer client library will be used.

  Returns:
    Operation: the requested operation
  """
    return GetService(release_track=release_track).Get(
        api_util.GetMessagesModule(release_track=release_track).
        ComposerProjectsLocationsOperationsGetRequest(
            name=operation_resource.RelativeName()))
def _ConstructCloudSqlMachineTypePatch(cloud_sql_machine_type, release_track):
    """Constructs an environment patch for Cloud SQL machine type.

  Args:
    cloud_sql_machine_type: str or None, Cloud SQL machine type used by the
      Airflow database.
    release_track: base.ReleaseTrack, the release track of command. It dictates
      which Composer client library is used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
    messages = api_util.GetMessagesModule(release_track=release_track)
    config = messages.EnvironmentConfig(databaseConfig=messages.DatabaseConfig(
        machineType=cloud_sql_machine_type))
    return 'config.database_config.machine_type', messages.Environment(
        config=config)
def RestartWebServer(environment_ref, release_track=base.ReleaseTrack.BETA):
    """Calls the Composer Environments.RestartWebServer method.

  Args:
    environment_ref: Resource, the Composer environment resource to restart the
      web server for.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
      which Composer client library will be used.

  Returns:
    Operation: the operation corresponding to the restart of the web server
  """
    message_module = api_util.GetMessagesModule(release_track=release_track)
    request_message = message_module.ComposerProjectsLocationsEnvironmentsRestartWebServerRequest(
        name=environment_ref.RelativeName())
    return GetService(
        release_track=release_track).RestartWebServer(request_message)
Exemple #30
0
def _ConstructImageVersionPatch(update_image_version,
                                release_track=base.ReleaseTrack.GA):
  """Constructs an environment patch for environment image version.

  Args:
    update_image_version: string, the target image version.
    release_track: base.ReleaseTrack, the release track of command. Will dictate
      which Composer client library will be used.

  Returns:
    (str, Environment), the field mask and environment to use for update.
  """
  messages = api_util.GetMessagesModule(release_track=release_track)
  software_config = messages.SoftwareConfig(imageVersion=update_image_version)
  config = messages.EnvironmentConfig(softwareConfig=software_config)

  return 'config.software_config.image_version', messages.Environment(
      config=config)