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 IsValidImageVersionUpgrade(env_ref,
image_version_id,
release_track=base.ReleaseTrack.GA):
"""Checks if image version candidate is a valid upgrade for environment."""
# Checks for the use of an alias and confirms that a valid airflow upgrade has
# been requested.
cand_image_ver = _ImageVersionItem(image_ver=image_version_id)
env_details = environments_api_util.Get(env_ref, release_track)
cur_image_ver = _ImageVersionItem(
image_ver=env_details.config.softwareConfig.imageVersion)
if not (CompareVersions(MIN_UPGRADEABLE_COMPOSER_VER,
cur_image_ver.composer_ver) <= 0):
raise InvalidImageVersionError(
'This environment does not support upgrades.')
if cand_image_ver.composer_contains_alias:
if _IsAirflowVersionUpgradeCompatible(cur_image_ver.airflow_ver,
cand_image_ver.airflow_ver):
return True
else:
# Leaves the validity check to the Composer backend request validation.
if _ValidateCandidateImageVersionId(
cur_image_ver.GetImageVersionString(),
cand_image_ver.GetImageVersionString()):
return True
return False
def ListImageVersionUpgrades(env_ref, release_track=base.ReleaseTrack.GA):
"""List of available image version upgrades for provided env_ref."""
env_details = environments_api_util.Get(env_ref, release_track)
proj_location_ref = env_ref.Parent()
cur_image_version_id = env_details.config.softwareConfig.imageVersion
cur_python_version = env_details.config.softwareConfig.pythonVersion
return _BuildUpgradeCandidateList(proj_location_ref, cur_image_version_id,
cur_python_version, release_track)
def _GetStorageBucket(env_ref, release_track=base.ReleaseTrack.GA):
env = environments_api_util.Get(env_ref, release_track=release_track)
if not env.config.dagGcsPrefix:
raise command_util.Error(BUCKET_MISSING_MSG)
try:
gcs_dag_dir = storage_util.ObjectReference.FromUrl(env.config.dagGcsPrefix)
except (storage_util.InvalidObjectNameError, ValueError):
raise command_util.Error(BUCKET_MISSING_MSG)
return gcs_dag_dir.bucket_ref
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 _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 Run(self, args):
self.DeprecationWarningPrompt(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):
try:
kubectl_ns = command_util.FetchKubectlNamespace(
env_obj.config.softwareConfig.imageVersion)
pod = command_util.GetGkePod(
pod_substr=WORKER_POD_SUBSTR, kubectl_namespace=kubectl_ns)
log.status.Print(
'Executing within the following kubectl namespace: {}'.format(
kubectl_ns))
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(
command_util.AddKubectlNamespace(kubectl_ns, kubectl_args),
out_func=log.status.Print)
except command_util.KubectlError as e:
raise self.ConvertKubectlError(e, env_obj)
def IsValidImageVersionUpgrade(env_ref,
image_version_id,
release_track=base.ReleaseTrack.GA):
"""Checks if image version candidate is a valid upgrade for environment."""
# Checks for the use of an alias and confirms that a valid airflow upgrade has
# been requested.
cand_image_ver = _ImageVersionItem(image_ver=image_version_id)
if cand_image_ver.composer_contains_alias:
env_details = environments_api_util.Get(env_ref, release_track)
cur_image_ver = _ImageVersionItem(
image_ver=env_details.config.softwareConfig.imageVersion)
if _IsAirflowVersionUpgradeCompatible(cur_image_ver.airflow_ver,
cand_image_ver.airflow_ver):
return True
else:
# Checks if supplied image_version_id matches an eligible env upgrade.
valid_image_versions = ListImageVersionUpgrades(env_ref, release_track)
return any(True for version in valid_image_versions
if version.imageVersionId.startswith(image_version_id))
def Run(self, args):
env_ref = args.CONCEPTS.environment.Parse()
env_obj = environments_api_util.Get(env_ref,
release_track=self.ReleaseTrack())
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
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))
kubectl_args = ['exec', pod, '--stdin']
if tty:
kubectl_args.append('--tty')
kubectl_args.extend(['--container', WORKER_CONTAINER, '--'])
if args.tree:
kubectl_args.extend(
['python', '-m', 'pipdeptree', '--warn'])
else:
kubectl_args.extend(['python', '-m', 'pip', 'list'])
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 ListImageVersionUpgrades(env_ref, release_track=base.ReleaseTrack.GA):
"""List of available image version upgrades for provided env_ref."""
env_details = environments_api_util.Get(env_ref, release_track)
cur_image_version_id = env_details.config.softwareConfig.imageVersion
cur_python_version = env_details.config.softwareConfig.pythonVersion
log.status.Print(
'Fetching list of available upgrades for image version \'{}\' ...'.
format(cur_image_version_id))
image_version_service = image_version_api_util.ImageVersionService(
release_track)
available_upgrades = []
for version in image_version_service.List(env_ref.Parent()):
try:
if (_ValidateCandidateVersionStrings(cur_image_version_id,
version.imageVersionId)
and cur_python_version in version.supportedPythonVersions):
available_upgrades.append(version)
except InvalidImageVersionError:
pass
return available_upgrades
def Run(self, args):
env_ref = args.CONCEPTS.environment.Parse()
return environments_api_util.Get(env_ref,
release_track=self.ReleaseTrack())
def _ConstructPatch(self,
env_ref,
args,
support_environment_upgrades=False):
env_obj = environments_api_util.Get(env_ref,
release_track=self.ReleaseTrack())
is_composer_v1 = image_versions_command_util.IsImageVersionStringComposerV1(
env_obj.config.softwareConfig.imageVersion)
params = dict(
is_composer_v1=is_composer_v1,
env_ref=env_ref,
node_count=args.node_count,
update_pypi_packages_from_file=args.update_pypi_packages_from_file,
clear_pypi_packages=args.clear_pypi_packages,
remove_pypi_packages=args.remove_pypi_packages,
update_pypi_packages=dict(
command_util.SplitRequirementSpecifier(r)
for r in args.update_pypi_package),
clear_labels=args.clear_labels,
remove_labels=args.remove_labels,
update_labels=args.update_labels,
clear_airflow_configs=args.clear_airflow_configs,
remove_airflow_configs=args.remove_airflow_configs,
update_airflow_configs=args.update_airflow_configs,
clear_env_variables=args.clear_env_variables,
remove_env_variables=args.remove_env_variables,
update_env_variables=args.update_env_variables,
release_track=self.ReleaseTrack())
if support_environment_upgrades:
params['update_image_version'] = args.image_version
params['update_web_server_access_control'] = (
environments_api_util.BuildWebServerAllowedIps(
args.update_web_server_allow_ip, args.web_server_allow_all,
args.web_server_deny_all))
if (args.cloud_sql_machine_type and not is_composer_v1):
raise command_util.InvalidUserInputError(
_INVALID_OPTION_FOR_V2_ERROR_MSG.format(
opt='cloud-sql-machine-type'))
if (args.web_server_machine_type and not is_composer_v1):
raise command_util.InvalidUserInputError(
_INVALID_OPTION_FOR_V2_ERROR_MSG.format(
opt='web-server-machine-type'))
params['cloud_sql_machine_type'] = args.cloud_sql_machine_type
params['web_server_machine_type'] = args.web_server_machine_type
params['scheduler_count'] = args.scheduler_count
if self._support_environment_size:
if (args.environment_size and is_composer_v1):
raise command_util.InvalidUserInputError(
_INVALID_OPTION_FOR_V1_ERROR_MSG.format(
opt='environment-size'))
if self.ReleaseTrack() == base.ReleaseTrack.GA:
params[
'environment_size'] = flags.ENVIRONMENT_SIZE_GA.GetEnumForChoice(
args.environment_size)
elif self.ReleaseTrack() == base.ReleaseTrack.BETA:
params[
'environment_size'] = flags.ENVIRONMENT_SIZE_BETA.GetEnumForChoice(
args.environment_size)
elif self.ReleaseTrack() == base.ReleaseTrack.ALPHA:
params[
'environment_size'] = flags.ENVIRONMENT_SIZE_ALPHA.GetEnumForChoice(
args.environment_size)
if self._support_autoscaling:
if (is_composer_v1 and
(args.scheduler_cpu or args.worker_cpu or args.web_server_cpu
or args.scheduler_memory or args.worker_memory
or args.web_server_memory or args.scheduler_storage
or args.worker_storage or args.web_server_storage
or args.min_workers or args.max_workers)):
raise command_util.InvalidUserInputError(
'Workloads Config flags introduced in Composer 2.X'
' cannot be used when creating Composer 1.X environments.')
params['scheduler_cpu'] = args.scheduler_cpu
params['worker_cpu'] = args.worker_cpu
params['web_server_cpu'] = args.web_server_cpu
params[
'scheduler_memory_gb'] = environments_api_util.MemorySizeBytesToGB(
args.scheduler_memory)
params[
'worker_memory_gb'] = environments_api_util.MemorySizeBytesToGB(
args.worker_memory)
params[
'web_server_memory_gb'] = environments_api_util.MemorySizeBytesToGB(
args.web_server_memory)
params[
'scheduler_storage_gb'] = environments_api_util.MemorySizeBytesToGB(
args.scheduler_storage)
params[
'worker_storage_gb'] = environments_api_util.MemorySizeBytesToGB(
args.worker_storage)
params[
'web_server_storage_gb'] = environments_api_util.MemorySizeBytesToGB(
args.web_server_storage)
params['min_workers'] = args.min_workers
params['max_workers'] = args.max_workers
if self._support_maintenance_window:
params['maintenance_window_start'] = args.maintenance_window_start
params['maintenance_window_end'] = args.maintenance_window_end
params[
'maintenance_window_recurrence'] = args.maintenance_window_recurrence
if self.ReleaseTrack() != base.ReleaseTrack.GA:
if args.enable_master_authorized_networks and args.disable_master_authorized_networks:
raise command_util.InvalidUserInputError(
'Cannot specify --enable-master-authorized-networks with --disable-master-authorized-networks'
)
if args.disable_master_authorized_networks and args.master_authorized_networks:
raise command_util.InvalidUserInputError(
'Cannot specify --disable-master-authorized-networks with --master-authorized-networks'
)
if args.enable_master_authorized_networks is None and args.master_authorized_networks:
raise command_util.InvalidUserInputError(
'Cannot specify --master-authorized-networks without --enable-master-authorized-networks'
)
if args.enable_master_authorized_networks or args.disable_master_authorized_networks:
params[
'master_authorized_networks_enabled'] = True if args.enable_master_authorized_networks else False
command_util.ValidateMasterAuthorizedNetworks(
args.master_authorized_networks)
params[
'master_authorized_networks'] = args.master_authorized_networks
return patch_util.ConstructPatch(**params)
def Run(self, args):
env_ref = args.CONCEPTS.environment.Parse()
return environments_api_util.Get(env_ref)
