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')
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)
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)
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)
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
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()))
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)
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)