362 lines
9.4 KiB
Python
362 lines
9.4 KiB
Python
# -*- coding: utf-8 -*- #
|
|
# Copyright 2024 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.
|
|
"""Flags and helpers for the Managed Flink CLI."""
|
|
|
|
import argparse
|
|
from googlecloudsdk.calliope import arg_parsers
|
|
from googlecloudsdk.calliope import base
|
|
from googlecloudsdk.command_lib.util import completers
|
|
from googlecloudsdk.command_lib.util import parameter_info_lib
|
|
|
|
_AUTOTUNING_MODES = {
|
|
'fixed': 'The number of taskmanagers is fixed.',
|
|
'elastic': (
|
|
'The number of taskmanagers is scaled automatically according to'
|
|
' workload.'
|
|
),
|
|
}
|
|
|
|
# Completers
|
|
|
|
|
|
class ListCommandParameterInfo(parameter_info_lib.ParameterInfoByConvention):
|
|
"""Helper class for ListCommandCompleter."""
|
|
|
|
def GetFlag(
|
|
self,
|
|
parameter_name,
|
|
parameter_value=None,
|
|
check_properties=True,
|
|
for_update=False,
|
|
):
|
|
return super(ListCommandParameterInfo, self).GetFlag(
|
|
parameter_name,
|
|
parameter_value,
|
|
check_properties=check_properties,
|
|
for_update=for_update,
|
|
)
|
|
|
|
|
|
class ListCommandCompleter(completers.ListCommandCompleter):
|
|
"""Helper class for DeploymentCompleter."""
|
|
|
|
def ParameterInfo(self, parsed_args, arguments):
|
|
return ListCommandParameterInfo(
|
|
parsed_args,
|
|
arguments,
|
|
self.collection,
|
|
updaters=COMPLETERS_BY_CONVENTION,
|
|
)
|
|
|
|
|
|
class DeploymentCompleter(ListCommandCompleter):
|
|
"""Completer for listing deployments."""
|
|
|
|
def __init__(self, **kwargs):
|
|
super(DeploymentCompleter, self).__init__(
|
|
collection='managedflink.projects.locations.deployments',
|
|
list_command='managed-flink deployments list',
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
COMPLETERS_BY_CONVENTION = {'deployment': (DeploymentCompleter, False)}
|
|
|
|
|
|
# Flags
|
|
def AddNetworkConfigArgs(parser):
|
|
"""Adds network config arguments."""
|
|
parser.add_argument(
|
|
'--network-config-vpc',
|
|
metavar='NETWORK',
|
|
dest='network',
|
|
help='The network to use for the job.',
|
|
)
|
|
parser.add_argument(
|
|
'--network-config-subnetwork',
|
|
metavar='SUBNETWORK',
|
|
dest='subnetwork',
|
|
help='The subnetwork to use for the job.',
|
|
)
|
|
|
|
|
|
def AddWorkloadIdentityArgument(parser):
|
|
"""Adds workload identity argument."""
|
|
parser.add_argument(
|
|
'--workload-identity',
|
|
metavar='WORKLOAD_IDENTITY',
|
|
dest='workload_identity',
|
|
help=(
|
|
'The workload identity to use for the job. Managed Flink'
|
|
' Default Workload Identity will be used if not specified.'
|
|
),
|
|
)
|
|
|
|
|
|
def AddLocationArgument(parser):
|
|
"""Creates location argument."""
|
|
base.Argument(
|
|
'--location',
|
|
metavar='LOCATION',
|
|
required=True,
|
|
dest='location',
|
|
suggestion_aliases=['--region'],
|
|
help='The location to run the job in.',
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddJobTypeArgument(parser):
|
|
"""Job type arguments."""
|
|
base.Argument(
|
|
'--job-type',
|
|
metavar='JOB_TYPE',
|
|
choices=['auto', 'jar', 'python', 'sql'],
|
|
default='auto',
|
|
help=(
|
|
'The type of job to run. If "auto" will be selected based on the file'
|
|
' extension for the job argument.'
|
|
),
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddJobJarArgument(parser):
|
|
"""Creates the job argument."""
|
|
base.Argument(
|
|
'job',
|
|
metavar='JAR|PY|SQL',
|
|
help=(
|
|
'The file containing the Flink job to run. Can be a jar, python, or'
|
|
' sql file.'
|
|
),
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddExtraJarsArgument(parser):
|
|
"""Creates the extra jars argument."""
|
|
base.Argument(
|
|
'--jars',
|
|
metavar='JAR',
|
|
type=arg_parsers.ArgList(),
|
|
dest='extra_jars',
|
|
help=(
|
|
'The extra jars to pass to the job. Can be a jar, python, or'
|
|
' sql file.'
|
|
),
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddDryRunArgument(parser):
|
|
"""Creates dry run argument."""
|
|
base.Argument(
|
|
'--dry-run',
|
|
action='store_true',
|
|
dest='dry_run',
|
|
default=False,
|
|
required=False,
|
|
help='Return command used to submit a job without invoking API.',
|
|
).AddToParser(parser)
|
|
|
|
|
|
# This has been temporarily disabled. Commented out to avoid confusing
|
|
# test coverage.
|
|
# def AddManagedKafkaClustersArgument(parser):
|
|
# """Creates the managed flink argument."""
|
|
# base.Argument(
|
|
# '--managed-kafka-clusters',
|
|
# metavar='MANAGED_KAFKA_CLUSTERS',
|
|
# dest='managed_kafka_clusters',
|
|
# type=arg_parsers.ArgList(),
|
|
# help='Specifies managed kafka clusters to associate with this job.',
|
|
# ).AddToParser(parser)
|
|
|
|
|
|
def AddMainClassArgument(parser):
|
|
"""Creates main class argument."""
|
|
base.Argument(
|
|
'--class',
|
|
metavar='CLASS',
|
|
dest='main_class',
|
|
help=(
|
|
'The main class of the Flink job. Required if the jar file manifest'
|
|
' does not contain a main class.'
|
|
),
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddJobArgsCollector(parser):
|
|
"""Collects extra arguments into the job_args list."""
|
|
parser.add_argument(
|
|
'job_args',
|
|
nargs=argparse.REMAINDER,
|
|
help='The job arguments to pass.',
|
|
)
|
|
|
|
|
|
def AddNameArgument(parser):
|
|
"""Creates name argument."""
|
|
base.Argument(
|
|
'--name',
|
|
metavar='NAME',
|
|
dest='name',
|
|
required=False,
|
|
help='The name of the job. The Flink job name will be used if not set.',
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddJobIdArgument(parser):
|
|
"""Creates job id argument."""
|
|
base.Argument(
|
|
'job_id',
|
|
metavar='JOBID',
|
|
help='The id of the job.',
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddAsyncArgument(parser, default=False):
|
|
"""Creates async argument."""
|
|
base.Argument(
|
|
'--async',
|
|
action='store_true',
|
|
dest='async_submit',
|
|
default=default,
|
|
required=False,
|
|
help='Return immediately after job submission.',
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddStagingLocationArgument(parser):
|
|
"""Creates staging location argument."""
|
|
base.Argument(
|
|
'--staging-location',
|
|
metavar='STAGING_LOCATION',
|
|
dest='staging_location',
|
|
required=True,
|
|
help=(
|
|
'The Google Cloud Storage staging location for the job. Must start'
|
|
' with gs://'
|
|
),
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddDeploymentArgument(
|
|
parser,
|
|
help_text_to_prepend=None,
|
|
help_text_to_overwrite=None,
|
|
required=False,
|
|
):
|
|
"""Creates deployment argument."""
|
|
|
|
if help_text_to_overwrite:
|
|
help_text = help_text_to_overwrite
|
|
else:
|
|
help_text = """
|
|
The Flink Deployment to use for this invocation.
|
|
"""
|
|
|
|
if help_text_to_prepend:
|
|
help_text = '\n\n'.join((help_text_to_prepend, help_text))
|
|
|
|
base.Argument(
|
|
'--deployment',
|
|
metavar='DEPLOYMENT_NAME',
|
|
required=required,
|
|
dest='deployment',
|
|
completer=DeploymentCompleter,
|
|
help=help_text,
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddAutotuningModeArgument(parser, default='elastic', required=False):
|
|
"""Creates autotuning mode argument."""
|
|
|
|
base.Argument(
|
|
'--autotuning-mode',
|
|
metavar='AUTOTUNING_MODE',
|
|
choices=_AUTOTUNING_MODES,
|
|
default=default,
|
|
required=required,
|
|
dest='autotuning_mode',
|
|
help='Selects the autotuning mode for jobs.',
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddFixedParallelismArgs(parser):
|
|
"""Adds fixed parallelism arguments."""
|
|
parser.add_argument(
|
|
'--parallelism',
|
|
type=arg_parsers.BoundedInt(lower_bound=1, upper_bound=10000),
|
|
help='The parallelism of the job when in "fixed" autotuning mode.',
|
|
)
|
|
|
|
|
|
def AddElasticParallelismArgs(parser):
|
|
"""Adds elastic parallelism arguments."""
|
|
parser.add_argument(
|
|
'--min-parallelism',
|
|
type=arg_parsers.BoundedInt(lower_bound=1, upper_bound=10000),
|
|
help=(
|
|
'The minimum parallelism of the job when in "elastic" autotuning'
|
|
' mode. This will also be the initial parallelism of the job.'
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
'--max-parallelism',
|
|
type=arg_parsers.BoundedInt(lower_bound=1, upper_bound=10000),
|
|
help=(
|
|
'The maximum parallelism of the job when in "elastic" autotuning'
|
|
' mode.'
|
|
),
|
|
)
|
|
|
|
|
|
def AddShowOutputArgument(parser):
|
|
"""Creates show output argument."""
|
|
base.Argument(
|
|
'--enable-output',
|
|
action='store_true',
|
|
dest='show_output',
|
|
default=False,
|
|
required=False,
|
|
help='Shows the output of the Flink client.',
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddExtraArchivesArgument(parser):
|
|
"""Creates the extra archives argument."""
|
|
base.Argument(
|
|
'--archives',
|
|
metavar='ZIP',
|
|
type=arg_parsers.ArgList(),
|
|
dest='archives',
|
|
help=(
|
|
'The extra archives to pass to the job. Can be a zip file containing'
|
|
' resource files for the job.'
|
|
),
|
|
).AddToParser(parser)
|
|
|
|
|
|
def AddPythonVirtualEnvArgument(parser):
|
|
"""Creates main class argument."""
|
|
base.Argument(
|
|
'--python-venv',
|
|
metavar='ZIP',
|
|
dest='python_venv',
|
|
help=(
|
|
'The path to the zip file to manage the virtualenv for Python'
|
|
' dependencies. Required if the job type is python. Must start with'
|
|
' gs://.'
|
|
),
|
|
).AddToParser(parser)
|