349 lines
13 KiB
Python
349 lines
13 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2025 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.
|
|
"""Command to migrate an environment."""
|
|
|
|
from googlecloudsdk.api_lib.composer import environments_util as environments_api_util
|
|
from googlecloudsdk.api_lib.composer import operations_util as operations_api_util
|
|
from googlecloudsdk.api_lib.composer import util as api_util
|
|
from googlecloudsdk.calliope import arg_parsers
|
|
from googlecloudsdk.calliope import base
|
|
from googlecloudsdk.command_lib.composer import flags
|
|
from googlecloudsdk.command_lib.composer import image_versions_util as image_versions_command_util
|
|
from googlecloudsdk.command_lib.composer import resource_args
|
|
from googlecloudsdk.command_lib.composer import util as command_util
|
|
from googlecloudsdk.core import log
|
|
import six
|
|
|
|
|
|
@base.DefaultUniverseOnly
|
|
@base.Hidden
|
|
@base.ReleaseTracks(base.ReleaseTrack.ALPHA, base.ReleaseTrack.BETA)
|
|
class Migrate(base.Command):
|
|
"""Migrates an environment from Composer 2 to Composer 3 in-place.
|
|
|
|
If run asynchronously with `--async`, exits after printing an operation
|
|
that can be used to poll the status of the creation operation via:
|
|
|
|
{top_command} composer operations describe
|
|
"""
|
|
|
|
@classmethod
|
|
def Args(cls, parser):
|
|
resource_args.AddEnvironmentResourceArg(parser, 'to migrate')
|
|
base.ASYNC_FLAG.AddToParser(parser)
|
|
|
|
target_version_type = arg_parsers.RegexpValidator(
|
|
r'^composer-3-airflow-(\d+(?:\.\d+(?:\.\d+(?:-build\.\d+)?)?)?)',
|
|
"must be in the form 'composer-3-airflow-X[.Y[.Z]]' For example:"
|
|
" 'composer-3-airflow-2.3.4-build.5'. Only migrations to Composer 3 are"
|
|
" supported.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--image-version',
|
|
type=target_version_type,
|
|
metavar='IMAGE_VERSION',
|
|
required=True,
|
|
help="""\
|
|
Migrate the Composer 2 environment to this Composer 3 version in-place.
|
|
|
|
The image version encapsulates the versions of both Cloud Composer and
|
|
Apache Airflow. Must be of the form
|
|
`composer-3-airflow-X[.Y[.Z[-build.N]]]`, where `[]` denotes optional
|
|
fragments.
|
|
|
|
Examples: `composer-3-airflow-2`, `composer-3-airflow-2.2`,
|
|
`composer-3-airflow-2.3.4`, `composer-3-airflow-2.3.4-build.5`.
|
|
|
|
The Cloud Composer portion of the image version must be a Composer 3
|
|
version. The Apache Airflow portion of the image version can be a
|
|
semantic version or an alias in the form of major or major.minor
|
|
version numbers, resolved to the latest matching Apache Airflow version
|
|
supported in the given Cloud Composer version. The resolved versions are
|
|
stored in the migrated environment.
|
|
""",
|
|
)
|
|
|
|
gke_cluster_retention_policy_group = parser.add_mutually_exclusive_group(
|
|
help=(
|
|
'Specify what should happen to the Composer 2 GKE cluster after'
|
|
' migration. If cluster is retained, until deleted manually it will'
|
|
' contribute to enviornment\'s cost.'
|
|
),
|
|
required=True,
|
|
)
|
|
gke_cluster_retention_policy_group.add_argument(
|
|
'--retain-gke-cluster',
|
|
action='store_true',
|
|
help='Retain Composer 2 GKE cluster after migration.',
|
|
)
|
|
gke_cluster_retention_policy_group.add_argument(
|
|
'--delete-gke-cluster',
|
|
action='store_true',
|
|
help='Delete Composer 2 GKE cluster after migration.',
|
|
)
|
|
|
|
maintenance_window_group = parser.add_argument_group(
|
|
help=(
|
|
'Specify the maintenance window for the migrated environment. It'
|
|
' will override the current maintenance window. If not specified,'
|
|
' and the enviornment uses Composer 2 default maintenance window,'
|
|
' the migrated environment will use the Composer 3 default'
|
|
' maintenance window.'
|
|
)
|
|
)
|
|
flags.MAINTENANCE_WINDOW_START_FLAG.AddToParser(maintenance_window_group)
|
|
flags.MAINTENANCE_WINDOW_END_FLAG.AddToParser(maintenance_window_group)
|
|
flags.MAINTENANCE_WINDOW_RECURRENCE_FLAG.AddToParser(
|
|
maintenance_window_group
|
|
)
|
|
|
|
dag_processor_group = parser.add_argument_group(
|
|
required=True,
|
|
help=(
|
|
'Specify the configuration of DAG processor for the'
|
|
' Composer 3 environment.'
|
|
)
|
|
)
|
|
dag_processor_group.add_argument(
|
|
'--dag-processor-cpu',
|
|
type=float,
|
|
required=True,
|
|
help='CPU allocated to Airflow dag processor'
|
|
)
|
|
dag_processor_group.add_argument(
|
|
'--dag-processor-memory',
|
|
type=arg_parsers.BinarySize(
|
|
lower_bound='1GB',
|
|
upper_bound='128GB',
|
|
suggested_binary_size_scales=['MB', 'GB'],
|
|
default_unit='G',
|
|
),
|
|
required=True,
|
|
help=(
|
|
'Memory allocated to Airflow dag processor, ex. 1GB, 3GB, 2. If'
|
|
' units are not provided, defaults to GB.'
|
|
),
|
|
)
|
|
dag_processor_group.add_argument(
|
|
'--dag-processor-storage',
|
|
type=arg_parsers.BinarySize(
|
|
default_unit='G',
|
|
),
|
|
required=True,
|
|
help=(
|
|
'Storage allocated to Airflow dag processor, ex. 600MB, 3GB, 2. If'
|
|
' units are not provided, defaults to GB.'
|
|
),
|
|
)
|
|
dag_processor_group.add_argument(
|
|
'--dag-processor-count',
|
|
type=int,
|
|
required=True,
|
|
help='Number of dag processors',
|
|
)
|
|
|
|
scheduler_group = parser.add_argument_group(
|
|
help=(
|
|
'Group of arguments for setting scheduler configuration in migrated'
|
|
' Composer environment. If not specified, the current scheduler'
|
|
' configuration will be preserved.'
|
|
)
|
|
)
|
|
flags.SCHEDULER_CPU.AddToParser(scheduler_group)
|
|
flags.SCHEDULER_MEMORY.AddToParser(scheduler_group)
|
|
flags.SCHEDULER_STORAGE.AddToParser(scheduler_group)
|
|
flags.NUM_SCHEDULERS.AddToParser(scheduler_group)
|
|
|
|
worker_group = parser.add_argument_group(
|
|
help=(
|
|
'Group of arguments for setting worker configuration in migrated'
|
|
' Composer environment. If not specified, the current worker'
|
|
' configuration will be preserved.'
|
|
)
|
|
)
|
|
flags.WORKER_CPU.AddToParser(worker_group)
|
|
flags.WORKER_MEMORY.AddToParser(worker_group)
|
|
flags.WORKER_STORAGE.AddToParser(worker_group)
|
|
flags.MIN_WORKERS.AddToParser(worker_group)
|
|
flags.MAX_WORKERS.AddToParser(worker_group)
|
|
|
|
web_server_group = parser.add_argument_group(
|
|
help=(
|
|
'Group of arguments for setting web server configuration in'
|
|
' migrated Composer environment. If not specified, the current web'
|
|
' server configuration will be preserved.'
|
|
)
|
|
)
|
|
flags.WEB_SERVER_CPU.AddToParser(web_server_group)
|
|
flags.WEB_SERVER_MEMORY.AddToParser(web_server_group)
|
|
flags.WEB_SERVER_STORAGE.AddToParser(web_server_group)
|
|
|
|
triggerer_group = parser.add_argument_group(
|
|
help=(
|
|
'Group of arguments for setting triggerer configuration in migrated'
|
|
' Composer environment. If not specified, the current triggerer'
|
|
' configuration will be preserved.'
|
|
)
|
|
)
|
|
flags.TRIGGERER_CPU.AddToParser(triggerer_group)
|
|
flags.TRIGGERER_MEMORY.AddToParser(triggerer_group)
|
|
flags.TRIGGERER_COUNT.AddToParser(triggerer_group)
|
|
|
|
def _Validate(self, env_obj, args):
|
|
if not image_versions_command_util.IsImageVersionStringComposerV2(
|
|
env_obj.config.softwareConfig.imageVersion
|
|
) or not image_versions_command_util.IsImageVersionStringComposerV3(
|
|
args.image_version
|
|
):
|
|
raise command_util.InvalidUserInputError(
|
|
'Migration is only supported from Composer 2 to Composer 3.'
|
|
)
|
|
|
|
def _ConstructGkeClusterRetentionPolicy(self, args, release_track):
|
|
messages = api_util.GetMessagesModule(release_track=release_track)
|
|
if args.retain_gke_cluster:
|
|
return (
|
|
messages.MigrateEnvironmentRequest.GkeClusterRetentionPolicyValueValuesEnum.RETAIN_GKE_CLUSTER
|
|
)
|
|
elif args.delete_gke_cluster:
|
|
return (
|
|
messages.MigrateEnvironmentRequest.GkeClusterRetentionPolicyValueValuesEnum.DELETE_GKE_CLUSTER
|
|
)
|
|
else:
|
|
raise command_util.InvalidUserInputError(
|
|
'One of --retain-gke-cluster or --delete-gke-cluster must be'
|
|
' specified.'
|
|
)
|
|
|
|
def _ConstructMigrateEnvironmentRequest(self, args, release_track):
|
|
messages = api_util.GetMessagesModule(release_track=release_track)
|
|
workloads_config = dict(
|
|
dagProcessor=messages.DagProcessorResource(
|
|
cpu=args.dag_processor_cpu,
|
|
memoryGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.dag_processor_memory
|
|
),
|
|
storageGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.dag_processor_storage
|
|
),
|
|
count=args.dag_processor_count,
|
|
),
|
|
)
|
|
if (
|
|
args.scheduler_cpu
|
|
or args.scheduler_memory
|
|
or args.scheduler_storage
|
|
or args.scheduler_count
|
|
):
|
|
workloads_config['scheduler'] = messages.SchedulerResource(
|
|
cpu=args.scheduler_cpu,
|
|
memoryGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.scheduler_memory
|
|
),
|
|
storageGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.scheduler_storage
|
|
),
|
|
count=args.scheduler_count,
|
|
)
|
|
if (
|
|
args.worker_cpu
|
|
or args.worker_memory
|
|
or args.worker_storage
|
|
or args.min_workers
|
|
or args.max_workers
|
|
):
|
|
workloads_config['worker'] = messages.WorkerResource(
|
|
cpu=args.worker_cpu,
|
|
memoryGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.worker_memory
|
|
),
|
|
storageGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.worker_storage
|
|
),
|
|
minCount=args.min_workers,
|
|
maxCount=args.max_workers,
|
|
)
|
|
if args.web_server_cpu or args.web_server_memory or args.web_server_storage:
|
|
workloads_config['webServer'] = messages.WebServerResource(
|
|
cpu=args.web_server_cpu,
|
|
memoryGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.web_server_memory
|
|
),
|
|
storageGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.web_server_storage
|
|
),
|
|
)
|
|
if args.triggerer_cpu or args.triggerer_memory or args.triggerer_count:
|
|
workloads_config['triggerer'] = messages.TriggererResource(
|
|
cpu=args.triggerer_cpu,
|
|
memoryGb=environments_api_util.MemorySizeBytesToGB(
|
|
args.triggerer_memory
|
|
),
|
|
count=args.triggerer_count,
|
|
)
|
|
|
|
migrate_request = dict(
|
|
imageVersion=args.image_version,
|
|
workloadsConfig=messages.WorkloadsConfig(**workloads_config),
|
|
gkeClusterRetentionPolicy=self._ConstructGkeClusterRetentionPolicy(
|
|
args, release_track
|
|
),
|
|
)
|
|
if (
|
|
args.maintenance_window_start
|
|
and args.maintenance_window_end
|
|
and args.maintenance_window_recurrence
|
|
):
|
|
migrate_request['maintenanceWindow'] = messages.MaintenanceWindow(
|
|
startTime=args.maintenance_window_start.isoformat(),
|
|
endTime=args.maintenance_window_end.isoformat(),
|
|
recurrence=args.maintenance_window_recurrence,
|
|
)
|
|
return messages.MigrateEnvironmentRequest(**migrate_request)
|
|
|
|
def Run(self, args):
|
|
env_ref = args.CONCEPTS.environment.Parse()
|
|
env_obj = environments_api_util.Get(
|
|
env_ref, release_track=self.ReleaseTrack())
|
|
self._Validate(env_obj, args)
|
|
|
|
request = self._ConstructMigrateEnvironmentRequest(
|
|
args, self.ReleaseTrack()
|
|
)
|
|
operation = environments_api_util.Migrate(
|
|
environment_ref=env_ref,
|
|
request=request,
|
|
release_track=self.ReleaseTrack(),
|
|
)
|
|
|
|
if args.async_:
|
|
log.UpdatedResource(
|
|
env_ref.RelativeName(),
|
|
kind='environment',
|
|
is_async=True,
|
|
details='with operation [{0}]'.format(operation.name),
|
|
)
|
|
return operation
|
|
try:
|
|
operations_api_util.WaitForOperation(
|
|
operation,
|
|
'Waiting for [{}] to be updated with [{}]'.format(
|
|
env_ref.RelativeName(), operation.name),
|
|
release_track=self.ReleaseTrack())
|
|
except command_util.Error as e:
|
|
raise command_util.Error('Error updating [{}]: {}'.format(
|
|
env_ref.RelativeName(), six.text_type(e)))
|