# -*- coding: utf-8 -*- #
# Copyright 2015 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.
"""Helpers for interacting with the Cloud Dataflow API."""


import json
import os
import shutil
import stat
import textwrap

from apitools.base.py import encoding
from apitools.base.py import exceptions as apitools_exceptions
from googlecloudsdk.api_lib.cloudbuild import cloudbuild_util
from googlecloudsdk.api_lib.storage import storage_api
from googlecloudsdk.api_lib.storage import storage_util
from googlecloudsdk.api_lib.util import apis
from googlecloudsdk.api_lib.util import exceptions
from googlecloudsdk.command_lib.builds import submit_util
from googlecloudsdk.core import log
from googlecloudsdk.core import properties
from googlecloudsdk.core.util import files
from googlecloudsdk.generated_clients.apis.dataflow.v1b3 import dataflow_v1b3_messages as messages
import six

DATAFLOW_API_NAME = 'dataflow'
DATAFLOW_API_VERSION = 'v1b3'
# TODO(b/139889563): Remove when dataflow args region is changed to required
DATAFLOW_API_DEFAULT_REGION = 'us-central1'


def GetMessagesModule():
  return apis.GetMessagesModule(DATAFLOW_API_NAME, DATAFLOW_API_VERSION)


def GetClientInstance():
  return apis.GetClientInstance(DATAFLOW_API_NAME, DATAFLOW_API_VERSION)


def GetProject():
  return properties.VALUES.core.project.Get(required=True)


def _GetBaseImagePath(image, is_distroless=False):
  """Returns full the image path of the given image."""
  if not is_distroless:
    return (
        f'gcr.io/dataflow-templates-base/{image}-template-launcher-base:latest'
    )
  return f'gcr.io/dataflow-templates-base/{image}-template-launcher-base-distroless:latest'


class Jobs:
  """The Jobs set of Dataflow API functions."""

  GET_REQUEST = GetMessagesModule().DataflowProjectsLocationsJobsGetRequest
  LIST_REQUEST = GetMessagesModule().DataflowProjectsLocationsJobsListRequest
  AGGREGATED_LIST_REQUEST = GetMessagesModule(
  ).DataflowProjectsJobsAggregatedRequest
  UPDATE_REQUEST = GetMessagesModule(
  ).DataflowProjectsLocationsJobsUpdateRequest

  @staticmethod
  def GetService():
    return GetClientInstance().projects_locations_jobs

  @staticmethod
  def Get(job_id, project_id=None, region_id=None, view=None):
    """Calls the Dataflow Jobs.Get method.

    Args:
      job_id: Identifies a single job.
      project_id: The project which owns the job.
      region_id: The regional endpoint where the job lives.
      view: (DataflowProjectsJobsGetRequest.ViewValueValuesEnum) Level of
        information requested in response.

    Returns:
      (Job)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    request = GetMessagesModule().DataflowProjectsLocationsJobsGetRequest(
        jobId=job_id, location=region_id, projectId=project_id, view=view)
    try:
      return Jobs.GetService().Get(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def Cancel(job_id, force=False, project_id=None, region_id=None):
    """Cancels a job by calling the Jobs.Update method.

    Args:
      job_id: Identifies a single job.
      force: True to forcibly cancel the job.
      project_id: The project which owns the job.
      region_id: The regional endpoint where the job lives.

    Returns:
      (Job)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    labels = None
    if force:
      labels = GetMessagesModule().Job.LabelsValue(additionalProperties=[
          GetMessagesModule().Job.LabelsValue.AdditionalProperty(
              key='force_cancel_job', value='true')
      ])
    job = GetMessagesModule().Job(
        labels=labels,
        requestedState=(GetMessagesModule().Job.RequestedStateValueValuesEnum
                        .JOB_STATE_CANCELLED))
    request = GetMessagesModule().DataflowProjectsLocationsJobsUpdateRequest(
        jobId=job_id, location=region_id, projectId=project_id, job=job)
    try:
      return Jobs.GetService().Update(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def UpdateOptions(
      job_id,
      project_id=None,
      region_id=None,
      min_num_workers=None,
      max_num_workers=None,
      worker_utilization_hint=None,
      unset_worker_utilization_hint=None,
      latency_tier=None,
      unset_latency_tier=None,
  ):
    """Update pipeline options on a running job.

    You should specify at-least one (or both) of min_num_workers and
    max_num_workers.

    Args:
      job_id: ID of job to update
      project_id: Project of the job
      region_id: Region the job is in
      min_num_workers: Lower-bound for worker autoscaling
      max_num_workers: Upper-bound for worker autoscaling
      worker_utilization_hint: Target CPU utilization for worker autoscaling
      unset_worker_utilization_hint: Unsets worker_utilization_hint value
      latency_tier: The latency tier for autoscaling
      unset_latency_tier: Unsets latency_tier value

    Returns:
      The updated Job
    """

    project_id = project_id or GetProject()
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    job = GetMessagesModule().Job(
        runtimeUpdatableParams=GetMessagesModule().RuntimeUpdatableParams(
            minNumWorkers=min_num_workers,
            maxNumWorkers=max_num_workers,
            workerUtilizationHint=(
                None
                if unset_worker_utilization_hint
                else worker_utilization_hint
            ),
            latencyTier=(None if unset_latency_tier else latency_tier),
        )
    )

    update_mask_pieces = []
    if min_num_workers is not None:
      update_mask_pieces.append('runtime_updatable_params.min_num_workers')
    if max_num_workers is not None:
      update_mask_pieces.append('runtime_updatable_params.max_num_workers')
    if (
        worker_utilization_hint is not None
        or unset_worker_utilization_hint
    ):
      update_mask_pieces.append(
          'runtime_updatable_params.worker_utilization_hint'
      )
    if latency_tier is not None or unset_latency_tier:
      update_mask_pieces.append('runtime_updatable_params.latency_tier')
    update_mask = ','.join(update_mask_pieces)

    request = GetMessagesModule().DataflowProjectsLocationsJobsUpdateRequest(
        jobId=job_id,
        location=region_id,
        projectId=project_id,
        job=job,
        updateMask=update_mask,
    )
    try:
      return Jobs.GetService().Update(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def Drain(job_id, project_id=None, region_id=None):
    """Drains a job by calling the Jobs.Update method.

    Args:
      job_id: Identifies a single job.
      project_id: The project which owns the job.
      region_id: The regional endpoint where the job lives.

    Returns:
      (Job)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    job = GetMessagesModule().Job(
        requestedState=(GetMessagesModule().Job.RequestedStateValueValuesEnum
                        .JOB_STATE_DRAINED))
    request = GetMessagesModule().DataflowProjectsLocationsJobsUpdateRequest(
        jobId=job_id, location=region_id, projectId=project_id, job=job)
    try:
      return Jobs.GetService().Update(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def Pause(
      job_id, project_id=None, region_id=None
  ) -> messages.Job:
    """Pauses a job by calling the Jobs.Update method.

    Args:
      job_id: Identifies a single job.
      project_id: The project which owns the job.
      region_id: The regional endpoint where the job lives.

    Returns:
      (messages.Job) The updated Job.
    """
    project_id = project_id or GetProject()
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    job = GetMessagesModule().Job(
        requestedState=(
            GetMessagesModule().Job.RequestedStateValueValuesEnum.JOB_STATE_PAUSING
        )
    )
    request = GetMessagesModule().DataflowProjectsLocationsJobsUpdateRequest(
        jobId=job_id, location=region_id, projectId=project_id, job=job
    )
    try:
      return Jobs.GetService().Update(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error) from error

  @staticmethod
  def Resume(
      job_id, project_id=None, region_id=None
  ) -> messages.Job:
    """Resumes a job by calling the Jobs.Update method.

    Args:
      job_id: Identifies a single job.
      project_id: The project which owns the job.
      region_id: The regional endpoint where the job lives.

    Returns:
      (Job)
    """
    project_id = project_id or GetProject()
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    job = GetMessagesModule().Job(
        requestedState=(
            GetMessagesModule().Job.RequestedStateValueValuesEnum.JOB_STATE_RUNNING
        )
    )
    request = GetMessagesModule().DataflowProjectsLocationsJobsUpdateRequest(
        jobId=job_id, location=region_id, projectId=project_id, job=job
    )
    try:
      return Jobs.GetService().Update(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error) from error

  @staticmethod
  def ResumeUnsupportedSDK(job_id,
                           experiment_with_token,
                           project_id=None,
                           region_id=None):
    """Resumes a job by calling the Jobs.Update method.

    Args:
      job_id: Identifies a single job.
      experiment_with_token: The resume token unique to the job prefixed with
        the experiment key.
      project_id: The project which owns the job.
      region_id: The regional endpoint where the job lives.

    Returns:
      (Job)
    """
    project_id = project_id or GetProject()
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    environment = GetMessagesModule().Environment(
        experiments=[experiment_with_token])
    job = GetMessagesModule().Job(environment=environment)
    request = GetMessagesModule().DataflowProjectsLocationsJobsUpdateRequest(
        jobId=job_id, location=region_id, projectId=project_id, job=job)
    try:
      return Jobs.GetService().Update(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def Snapshot(job_id,
               project_id=None,
               region_id=None,
               ttl='604800s',
               snapshot_sources=False):
    """Takes a snapshot of a job via the Jobs.Snapshot method.

    Args:
      job_id: Identifies a single job.
      project_id: The project which owns the job.
      region_id: The regional endpoint where the job lives.
      ttl: The ttl for the snapshot.
      snapshot_sources: If true, the sources will be snapshotted.

    Returns:
      (Snapshot)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    request = GetMessagesModule().DataflowProjectsLocationsJobsSnapshotRequest(
        jobId=job_id,
        location=region_id,
        projectId=project_id,
        snapshotJobRequest=GetMessagesModule().SnapshotJobRequest(
            location=region_id, ttl=ttl, snapshotSources=snapshot_sources),
    )
    try:
      return Jobs.GetService().Snapshot(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)


class Metrics:
  """The Metrics set of Dataflow API functions."""

  GET_REQUEST = GetMessagesModule(
  ).DataflowProjectsLocationsJobsGetMetricsRequest

  @staticmethod
  def GetService():
    return GetClientInstance().projects_locations_jobs

  @staticmethod
  def Get(job_id, project_id=None, region_id=None, start_time=None):
    """Calls the Dataflow Metrics.Get method.

    Args:
      job_id: The job to get messages for.
      project_id: The project which owns the job.
      region_id: The regional endpoint of the job.
      start_time: Return only metric data that has changed since this time.
        Default is to return all information about all metrics for the job.

    Returns:
      (MetricUpdate)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    request = GetMessagesModule(
    ).DataflowProjectsLocationsJobsGetMetricsRequest(
        jobId=job_id,
        location=region_id,
        projectId=project_id,
        startTime=start_time)
    try:
      return Metrics.GetService().GetMetrics(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)


class TemplateArguments:
  """Wrapper class for template arguments."""

  project_id = None
  region_id = None
  gcs_location = None
  job_name = None
  zone = None
  max_workers = None
  num_workers = None
  network = None
  subnetwork = None
  worker_machine_type = None
  launcher_machine_type = None
  launcher_vm_timeout_secs = None
  staging_location = None
  temp_location = None
  kms_key_name = None
  disable_public_ips = None
  parameters = None
  service_account_email = None
  worker_region = None
  worker_zone = None
  enable_streaming_engine = None
  additional_experiments = None
  additional_pipeline_options = None
  additional_user_labels = None
  streaming_update = None
  transform_name_mappings = None
  flexrs_goal = None

  def __init__(self,
               project_id=None,
               region_id=None,
               job_name=None,
               gcs_location=None,
               zone=None,
               max_workers=None,
               num_workers=None,
               network=None,
               subnetwork=None,
               worker_machine_type=None,
               launcher_machine_type=None,
               launcher_vm_timeout_secs=None,
               staging_location=None,
               temp_location=None,
               kms_key_name=None,
               disable_public_ips=None,
               parameters=None,
               service_account_email=None,
               worker_region=None,
               worker_zone=None,
               enable_streaming_engine=None,
               additional_experiments=None,
               additional_pipeline_options=None,
               additional_user_labels=None,
               streaming_update=None,
               transform_name_mappings=None,
               flexrs_goal=None):
    self.project_id = project_id
    self.region_id = region_id
    self.job_name = job_name
    self.gcs_location = gcs_location
    self.zone = zone
    self.max_workers = max_workers
    self.num_workers = num_workers
    self.network = network
    self.subnetwork = subnetwork
    self.worker_machine_type = worker_machine_type
    self.launcher_machine_type = launcher_machine_type
    self.launcher_vm_timeout_secs = launcher_vm_timeout_secs
    self.staging_location = staging_location
    self.temp_location = temp_location
    self.kms_key_name = kms_key_name
    self.disable_public_ips = disable_public_ips
    self.parameters = parameters
    self.service_account_email = service_account_email
    self.worker_region = worker_region
    self.worker_zone = worker_zone
    self.enable_streaming_engine = enable_streaming_engine
    self.additional_experiments = additional_experiments
    self.additional_pipeline_options = additional_pipeline_options
    self.additional_user_labels = additional_user_labels
    self.streaming_update = streaming_update
    self.transform_name_mappings = transform_name_mappings
    self.flexrs_goal = flexrs_goal


class Templates:
  """The Templates set of Dataflow API functions."""

  CREATE_REQUEST = GetMessagesModule().CreateJobFromTemplateRequest
  LAUNCH_TEMPLATE_PARAMETERS = GetMessagesModule().LaunchTemplateParameters
  LAUNCH_TEMPLATE_PARAMETERS_VALUE = LAUNCH_TEMPLATE_PARAMETERS.ParametersValue
  LAUNCH_FLEX_TEMPLATE_REQUEST = GetMessagesModule().LaunchFlexTemplateRequest
  PARAMETERS_VALUE = CREATE_REQUEST.ParametersValue
  FLEX_TEMPLATE_ENVIRONMENT = GetMessagesModule().FlexTemplateRuntimeEnvironment
  FLEX_TEMPLATE_USER_LABELS_VALUE = (
      FLEX_TEMPLATE_ENVIRONMENT.AdditionalUserLabelsValue
  )
  FLEX_TEMPLATE_PARAMETER = GetMessagesModule().LaunchFlexTemplateParameter
  FLEX_TEMPLATE_LAUNCH_OPTIONS_VALUE = (
      FLEX_TEMPLATE_PARAMETER.LaunchOptionsValue
  )
  DYNAMIC_TEMPLATE_TRANSFORM_NAME_MAPPING_VALUE = (
      LAUNCH_TEMPLATE_PARAMETERS.TransformNameMappingValue
  )
  FLEX_TEMPLATE_PARAMETER = GetMessagesModule().LaunchFlexTemplateParameter
  FLEX_TEMPLATE_PARAMETERS_VALUE = FLEX_TEMPLATE_PARAMETER.ParametersValue
  FLEX_TEMPLATE_TRANSFORM_NAME_MAPPING_VALUE = (
      FLEX_TEMPLATE_PARAMETER.TransformNameMappingsValue
  )
  IP_CONFIGURATION_ENUM_VALUE = GetMessagesModule(
  ).FlexTemplateRuntimeEnvironment.IpConfigurationValueValuesEnum
  FLEXRS_GOAL_ENUM_VALUE = GetMessagesModule(
  ).FlexTemplateRuntimeEnvironment.FlexrsGoalValueValuesEnum
  TEMPLATE_METADATA = GetMessagesModule().TemplateMetadata
  SDK_INFO = GetMessagesModule().SDKInfo
  SDK_LANGUAGE = GetMessagesModule().SDKInfo.LanguageValueValuesEnum
  CONTAINER_SPEC = GetMessagesModule().ContainerSpec
  FLEX_TEMPLATE_JAVA11_BASE_IMAGE = _GetBaseImagePath('java11')
  FLEX_TEMPLATE_JAVA17_BASE_IMAGE = _GetBaseImagePath('java17')
  FLEX_TEMPLATE_JAVA21_BASE_IMAGE = _GetBaseImagePath('java21')
  FLEX_TEMPLATE_JAVA25_BASE_IMAGE = _GetBaseImagePath('java25')
  FLEX_TEMPLATE_JAVA11_DISTROLESS_BASE_IMAGE = _GetBaseImagePath('java11', True)
  FLEX_TEMPLATE_JAVA17_DISTROLESS_BASE_IMAGE = _GetBaseImagePath('java17', True)
  FLEX_TEMPLATE_JAVA21_DISTROLESS_BASE_IMAGE = _GetBaseImagePath('java21', True)
  FLEX_TEMPLATE_JAVA25_DISTROLESS_BASE_IMAGE = _GetBaseImagePath('java25', True)
  FLEX_TEMPLATE_PYTHON3_BASE_IMAGE = _GetBaseImagePath('python3')
  FLEX_TEMPLATE_GO_BASE_IMAGE = _GetBaseImagePath('go')
  FLEX_TEMPLATE_GO_DISTROLESS_BASE_IMAGE = _GetBaseImagePath('go', True)
  YAML_TEMPLATE_GCS_LOCATION = (
      'gs://dataflow-templates-{}/latest/flex/Yaml_Template'
  )
  ALL_PERMISSIONS_MASK = (
      stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO
  )
  FILE_PERMISSIONS_MASK = (
      stat.S_IWUSR | stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH
  )
  # Directories need +x for access.
  DIR_PERMISSIONS_MASK = (
      stat.S_IRWXU | stat.S_IXGRP | stat.S_IRGRP | stat.S_IXOTH | stat.S_IROTH
  )

  @staticmethod
  def GetService():
    return GetClientInstance().projects_locations_templates

  @staticmethod
  def GetFlexTemplateService():
    return GetClientInstance().projects_locations_flexTemplates

  @staticmethod
  def Create(template_args=None):
    """Calls the Dataflow Templates.CreateFromJob method.

    Args:
      template_args: Arguments for create template.

    Returns:
      (Job)
    """
    params_list = []
    parameters = template_args.parameters
    for k, v in six.iteritems(parameters) if parameters else {}:
      params_list.append(
          Templates.PARAMETERS_VALUE.AdditionalProperty(key=k, value=v))

    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = template_args.region_id or DATAFLOW_API_DEFAULT_REGION

    ip_configuration_enum = GetMessagesModule(
    ).RuntimeEnvironment.IpConfigurationValueValuesEnum
    ip_private = ip_configuration_enum.WORKER_IP_PRIVATE
    ip_configuration = ip_private if template_args.disable_public_ips else None
    user_labels_value = GetMessagesModule(
    ).RuntimeEnvironment.AdditionalUserLabelsValue
    user_labels_list = Templates.__ConvertDictArguments(
        template_args.additional_user_labels,
        user_labels_value)

    body = Templates.CREATE_REQUEST(
        gcsPath=template_args.gcs_location,
        jobName=template_args.job_name,
        location=region_id,
        environment=GetMessagesModule().RuntimeEnvironment(
            serviceAccountEmail=template_args.service_account_email,
            zone=template_args.zone,
            maxWorkers=template_args.max_workers,
            numWorkers=template_args.num_workers,
            network=template_args.network,
            subnetwork=template_args.subnetwork,
            machineType=template_args.worker_machine_type,
            tempLocation=template_args.staging_location,
            kmsKeyName=template_args.kms_key_name,
            ipConfiguration=ip_configuration,
            workerRegion=template_args.worker_region,
            workerZone=template_args.worker_zone,
            enableStreamingEngine=template_args.enable_streaming_engine,
            additionalUserLabels=user_labels_value(
                additionalProperties=user_labels_list)
            if user_labels_list else None,
            additionalExperiments=(
                template_args.additional_experiments
                if template_args.additional_experiments else
                []
            ),
            additionalPipelineOptions=(
                template_args.additional_pipeline_options
                if template_args.additional_pipeline_options
                else []
            )),
        parameters=Templates.PARAMETERS_VALUE(
            additionalProperties=params_list) if parameters else None)
    request = GetMessagesModule(
    ).DataflowProjectsLocationsTemplatesCreateRequest(
        projectId=template_args.project_id or GetProject(),
        location=region_id,
        createJobFromTemplateRequest=body)

    try:
      return Templates.GetService().Create(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def LaunchDynamicTemplate(template_args=None):
    """Calls the Dataflow Templates.LaunchTemplate method on a dynamic template.

    Args:
      template_args: Arguments to create template. gcs_location must point to a
        Json serialized DynamicTemplateFileSpec.

    Returns:
      (LaunchTemplateResponse)
    """
    params_list = []
    parameters = template_args.parameters
    for k, v in six.iteritems(parameters) if parameters else {}:
      params_list.append(
          Templates.LAUNCH_TEMPLATE_PARAMETERS_VALUE.AdditionalProperty(
              key=k, value=v))

    transform_mapping_list = Templates.__ConvertDictArguments(
        template_args.transform_name_mappings,
        Templates.DYNAMIC_TEMPLATE_TRANSFORM_NAME_MAPPING_VALUE,
    )
    transform_mappings = None
    streaming_update = None
    if template_args.streaming_update:
      streaming_update = template_args.streaming_update
      if transform_mapping_list:
        transform_mappings = (
            Templates.DYNAMIC_TEMPLATE_TRANSFORM_NAME_MAPPING_VALUE(
                additionalProperties=transform_mapping_list
            )
        )

    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = template_args.region_id or DATAFLOW_API_DEFAULT_REGION

    ip_configuration_enum = GetMessagesModule(
    ).RuntimeEnvironment.IpConfigurationValueValuesEnum
    ip_private = ip_configuration_enum.WORKER_IP_PRIVATE
    ip_configuration = ip_private if template_args.disable_public_ips else None

    body = Templates.LAUNCH_TEMPLATE_PARAMETERS(
        environment=GetMessagesModule().RuntimeEnvironment(
            serviceAccountEmail=template_args.service_account_email,
            zone=template_args.zone,
            maxWorkers=template_args.max_workers,
            numWorkers=template_args.num_workers,
            network=template_args.network,
            subnetwork=template_args.subnetwork,
            machineType=template_args.worker_machine_type,
            tempLocation=template_args.staging_location,
            kmsKeyName=template_args.kms_key_name,
            ipConfiguration=ip_configuration,
            workerRegion=template_args.worker_region,
            workerZone=template_args.worker_zone,
            enableStreamingEngine=template_args.enable_streaming_engine,
            additionalExperiments=(
                template_args.additional_experiments
                if template_args.additional_experiments
                else []
            ),
        ),
        jobName=template_args.job_name,
        parameters=Templates.LAUNCH_TEMPLATE_PARAMETERS_VALUE(
            additionalProperties=params_list) if params_list else None,
        update=streaming_update,
        transformNameMapping=transform_mappings,
    )
    request = (
        GetMessagesModule().DataflowProjectsLocationsTemplatesLaunchRequest(
            gcsPath=template_args.gcs_location,
            location=region_id,
            launchTemplateParameters=body,
            projectId=template_args.project_id or GetProject(),
            validateOnly=False,
        )
    )

    try:
      return Templates.GetService().Launch(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def __ConvertDictArguments(arguments, value_message):
    """Convert dictionary arguments to parameter list .

    Args:
      arguments: Arguments for create job using template.
      value_message: the value message of the arguments

    Returns:
      List of value_message.AdditionalProperty
    """
    params_list = []
    if arguments:
      for k, v in six.iteritems(arguments):
        params_list.append(value_message.AdditionalProperty(key=k, value=v))

    return params_list

  @staticmethod
  def BuildJavaImageDockerfile(flex_template_base_image, pipeline_paths, env):
    """Builds Dockerfile contents for java flex template image.

    Args:
      flex_template_base_image: SDK version or base image to use.
      pipeline_paths: List of paths to pipelines and dependencies.
      env: Dictionary of env variables to set in the container image.

    Returns:
      Dockerfile contents as string.
    """
    dockerfile_template = """
    FROM {base_image}

    {env}

    {copy}

    {commands}
    """
    commands = ''
    env['FLEX_TEMPLATE_JAVA_CLASSPATH'] = '/template/*'
    envs = ['ENV {}={}'.format(k, v) for k, v in sorted(env.items())]
    env_list = '\n'.join(envs)
    paths = ' '.join(pipeline_paths)
    copy_command = 'COPY {} /template/'.format(paths)

    dockerfile_contents = textwrap.dedent(dockerfile_template).format(
        base_image=Templates._GetFlexTemplateBaseImage(
            flex_template_base_image),
        env=env_list,
        copy=copy_command,
        commands='\n'.join(commands))
    return dockerfile_contents

  @staticmethod
  def BuildPythonImageDockerfile(flex_template_base_image, pipeline_paths, env):
    """Builds Dockerfile contents for python flex template image.

    Args:
      flex_template_base_image: SDK version or base image to use.
      pipeline_paths: List of paths to pipelines and dependencies.
      env: Dictionary of env variables to set in the container image.

    Returns:
      Dockerfile contents as string.
    """
    dockerfile_template = """
    FROM {base_image}

    {env}

    {copy}

    {commands}
    """
    commands = [
        'apt-get update',
        'apt-get install -y libffi-dev git',
        'rm -rf /var/lib/apt/lists/*',
    ]

    env['FLEX_TEMPLATE_PYTHON_PY_FILE'] = (
        f'/template/{env["FLEX_TEMPLATE_PYTHON_PY_FILE"]}'
    )
    if 'FLEX_TEMPLATE_PYTHON_EXTRA_PACKAGES' in env:
      package_list = env['FLEX_TEMPLATE_PYTHON_EXTRA_PACKAGES'].split(',')
      if package_list:
        packages_path = [f'/template/{package}' for package in package_list]
        env['FLEX_TEMPLATE_PYTHON_EXTRA_PACKAGES'] = ','.join(packages_path)
        package_arg = ' '.join(packages_path)
        commands.append(
            f'pip install {package_arg}'
        )
    if 'FLEX_TEMPLATE_PYTHON_REQUIREMENTS_FILE' in env:
      env['FLEX_TEMPLATE_PYTHON_REQUIREMENTS_FILE'] = (
          f'/template/{env["FLEX_TEMPLATE_PYTHON_REQUIREMENTS_FILE"]}'
      )
      commands.append(
          'pip install --no-cache-dir -U -r'
          f' {env["FLEX_TEMPLATE_PYTHON_REQUIREMENTS_FILE"]}'
      )
      commands.append(
          '(pip check || (e=$?; echo "Building a container with incompatible'
          ' dependencies is prevented by default. If you are sure you want to'
          ' proceed, you need to create your own container image. See:'
          ' https://cloud.google.com/dataflow/docs/guides/templates/configuring-flex-templates";'
          ' exit $e))'
      )
    if 'FLEX_TEMPLATE_PYTHON_SETUP_FILE' in env:
      env['FLEX_TEMPLATE_PYTHON_SETUP_FILE'] = (
          f'/template/{env["FLEX_TEMPLATE_PYTHON_SETUP_FILE"]}'
      )

    envs = ['ENV {}={}'.format(k, v) for k, v in sorted(env.items())]
    env_list = '\n'.join(envs)
    paths = ' '.join(pipeline_paths)
    copy_command = 'COPY {} /template/'.format(paths)

    dockerfile_contents = textwrap.dedent(dockerfile_template).format(
        base_image=Templates._GetFlexTemplateBaseImage(
            flex_template_base_image),
        env=env_list,
        copy=copy_command,
        commands='RUN ' + ' && '.join(commands))
    return dockerfile_contents

  # staticmethod enforced by prior code, this violates the style guide and
  # would benefit from a refactor in the future.
  # TODO(b/242564654): Add type annotations for arguments when presubmits allow
  # them.
  @staticmethod
  def BuildGoImageDockerfile(flex_template_base_image,
                             pipeline_paths,
                             env):
    """Builds Dockerfile contents for go flex template image.

    Args:
      flex_template_base_image: SDK version or base image to use.
      pipeline_paths: Path to pipeline binary.
      env: Dictionary of env variables to set in the container image.

    Returns:
      Dockerfile contents as string.
    """
    dockerfile_template = """
    FROM {base_image}

    {env}

    {copy}
    """
    env['FLEX_TEMPLATE_GO_BINARY'] = '/template/{}'.format(
        env['FLEX_TEMPLATE_GO_BINARY'])
    paths = ' '.join(pipeline_paths)
    copy_command = 'COPY {} /template/'.format(paths)

    envs = [
        'ENV {}={}'.format(var, val) for var, val in sorted(env.items())
    ]
    env_list = '\n'.join(envs)

    dockerfile_contents = textwrap.dedent(dockerfile_template).format(
        base_image=Templates._GetFlexTemplateBaseImage(
            flex_template_base_image),
        env=env_list,
        copy=copy_command)

    return dockerfile_contents

  @staticmethod
  def BuildDockerfile(flex_template_base_image, pipeline_paths, env,
                      sdk_language):
    """Builds Dockerfile contents for flex template image.

    Args:
      flex_template_base_image: SDK version or base image to use.
      pipeline_paths: List of paths to pipelines and dependencies.
      env: Dictionary of env variables to set in the container image.
      sdk_language: SDK language of the flex template.

    Returns:
      Dockerfile contents as string.
    """
    if sdk_language == 'JAVA':
      return Templates.BuildJavaImageDockerfile(flex_template_base_image,
                                                pipeline_paths, env)
    elif sdk_language == 'PYTHON':
      return Templates.BuildPythonImageDockerfile(flex_template_base_image,
                                                  pipeline_paths, env)
    elif sdk_language == 'GO':
      return Templates.BuildGoImageDockerfile(flex_template_base_image,
                                              pipeline_paths, env)

  @staticmethod
  def _ValidateTemplateParameters(parameters):
    """Validates ParameterMetadata objects in template metadata.

    Args:
      parameters: List of ParameterMetadata objects.

    Raises:
      ValueError: If is any of the required field is not set.
    """
    for parameter in parameters:
      if not parameter.name:
        raise ValueError(
            'Invalid template metadata. Parameter name field is empty.'
            ' Parameter: {}'.format(parameter))
      if not parameter.label:
        raise ValueError(
            'Invalid template metadata. Parameter label field is empty.'
            ' Parameter: {}'.format(parameter))
      if not parameter.helpText:
        raise ValueError(
            'Invalid template metadata. Parameter helpText field is empty.'
            ' Parameter: {}'.format(parameter))

  @staticmethod
  def __ValidateFlexTemplateEnv(env, sdk_language):
    """Builds and validates Flex template environment values.

    Args:
      env: Dictionary of env variables to set in the container image.
      sdk_language: SDK language of the flex template.

    Returns:
      True on valid env values.

    Raises:
      ValueError: If is any of parameter value is invalid.
    """
    if sdk_language == 'JAVA' and 'FLEX_TEMPLATE_JAVA_MAIN_CLASS' not in env:
      raise ValueError(('FLEX_TEMPLATE_JAVA_MAIN_CLASS environment variable '
                        'should be provided for all JAVA jobs.'))
    elif sdk_language == 'PYTHON' and 'FLEX_TEMPLATE_PYTHON_PY_FILE' not in env:
      raise ValueError(('FLEX_TEMPLATE_PYTHON_PY_FILE environment variable '
                        'should be provided for all PYTHON jobs.'))
    elif sdk_language == 'GO' and 'FLEX_TEMPLATE_GO_BINARY' not in env:
      raise ValueError(('FLEX_TEMPLATE_GO_BINARY environment variable '
                        'should be provided for all GO jobs.'))
    return True

  @staticmethod
  def _BuildTemplateMetadata(template_metadata_json):
    """Builds and validates TemplateMetadata object.

    Args:
      template_metadata_json: Template metadata in json format.

    Returns:
      TemplateMetadata object on success.

    Raises:
      ValueError: If is any of the required field is not set.
    """
    template_metadata = encoding.JsonToMessage(Templates.TEMPLATE_METADATA,
                                               template_metadata_json)
    template_metadata_obj = Templates.TEMPLATE_METADATA()

    if not template_metadata.name:
      raise ValueError('Invalid template metadata. Name field is empty.'
                       ' Template Metadata: {}'.format(template_metadata))
    template_metadata_obj.name = template_metadata.name

    if template_metadata.description:
      template_metadata_obj.description = template_metadata.description

    if template_metadata.parameters:
      Templates._ValidateTemplateParameters(template_metadata.parameters)
      template_metadata_obj.parameters = template_metadata.parameters

    if template_metadata.yamlDefinition:
      template_metadata_obj.yamlDefinition = template_metadata.yamlDefinition

    return template_metadata_obj

  @staticmethod
  def GetYamlTemplateImage(args):
    """Returns the image path for a YAML template."""
    if args.image:
      return args.image
    elif args.yaml_image:
      return args.yaml_image

    # TODO: b/397983834 - Try to extract a region from the gcs bucket.
    if args.worker_region:
      try:
        return Templates._ExtractYamlTemplateImage(args.worker_region)
      except exceptions.HttpException:
        pass  # Fall through to using default region.

    return Templates._ExtractYamlTemplateImage(DATAFLOW_API_DEFAULT_REGION)

  @staticmethod
  def _ExtractYamlTemplateImage(region_id):
    """Returns the image path for a YAML template."""
    yaml_gcl_template_path = Templates.YAML_TEMPLATE_GCS_LOCATION.format(
        region_id
    )
    storage_client = storage_api.StorageClient()
    obj_ref = storage_util.ObjectReference.FromUrl(yaml_gcl_template_path)
    try:
      generic_template_definition = json.load(
          storage_client.ReadObject(obj_ref)
      )
    except Exception as e:
      raise exceptions.HttpException(
          'Unable to read file {0} due to incorrect file path or insufficient'
          ' read permissions'.format(yaml_gcl_template_path)
      ) from e
    return generic_template_definition['image']

  @staticmethod
  def _GetFlexTemplateBaseImage(flex_template_base_image):
    """Returns latest base image for given sdk version.

    Args:
        flex_template_base_image: SDK version or base image to use.

    Returns:
      If a custom base image value is given, returns the same value. Else,
      returns the latest base image for the given sdk version.
    """
    if flex_template_base_image == 'JAVA11':
      return Templates.FLEX_TEMPLATE_JAVA11_BASE_IMAGE
    elif flex_template_base_image == 'JAVA17':
      return Templates.FLEX_TEMPLATE_JAVA17_BASE_IMAGE
    elif flex_template_base_image == 'JAVA21':
      return Templates.FLEX_TEMPLATE_JAVA21_BASE_IMAGE
    elif flex_template_base_image == 'JAVA25':
      return Templates.FLEX_TEMPLATE_JAVA25_BASE_IMAGE
    elif flex_template_base_image == 'JAVA11_DISTROLESS':
      return Templates.FLEX_TEMPLATE_JAVA11_DISTROLESS_BASE_IMAGE
    elif flex_template_base_image == 'JAVA17_DISTROLESS':
      return Templates.FLEX_TEMPLATE_JAVA17_DISTROLESS_BASE_IMAGE
    elif flex_template_base_image == 'JAVA21_DISTROLESS':
      return Templates.FLEX_TEMPLATE_JAVA21_DISTROLESS_BASE_IMAGE
    elif flex_template_base_image == 'JAVA25_DISTROLESS':
      return Templates.FLEX_TEMPLATE_JAVA25_DISTROLESS_BASE_IMAGE
    elif flex_template_base_image == 'JAVA8':
      log.warning(
          'JAVA8 is deprecated and redirected to JAVA11. This option '
          'will be removed in a future release'
      )
      return Templates.FLEX_TEMPLATE_JAVA11_BASE_IMAGE
    elif flex_template_base_image == 'PYTHON3':
      return Templates.FLEX_TEMPLATE_PYTHON3_BASE_IMAGE
    elif flex_template_base_image == 'GO':
      return Templates.FLEX_TEMPLATE_GO_BASE_IMAGE
    elif flex_template_base_image == 'GO_DISTROLESS':
      return Templates.FLEX_TEMPLATE_GO_DISTROLESS_BASE_IMAGE
    return flex_template_base_image

  @staticmethod
  def _BuildSDKInfo(sdk_language):
    """Builds SDKInfo object.

    Args:
      sdk_language: SDK language of the flex template.

    Returns:
      SDKInfo object
    """
    if sdk_language == 'JAVA':
      return Templates.SDK_INFO(language=Templates.SDK_LANGUAGE.JAVA)
    elif sdk_language == 'PYTHON':
      return Templates.SDK_INFO(language=Templates.SDK_LANGUAGE.PYTHON)
    elif sdk_language == 'YAML':
      return Templates.SDK_INFO(language=Templates.SDK_LANGUAGE.YAML)
    elif sdk_language == 'GO':
      return Templates.SDK_INFO(language=Templates.SDK_LANGUAGE.GO)

  @staticmethod
  def _StoreFlexTemplateFile(template_file_gcs_location, container_spec_json):
    """Stores flex template container spec file in GCS.

    Args:
      template_file_gcs_location: GCS location to store the template file.
      container_spec_json: Container spec in json format.

    Returns:
      Returns the stored flex template file gcs object on success.
      Propagates the error on failures.
    """
    with files.TemporaryDirectory() as temp_dir:
      local_path = os.path.join(temp_dir, 'template-file.json')
      # Use Unix style line-endings on all platforms (especially Windows)
      files.WriteFileContents(local_path, container_spec_json, newline='\n')
      storage_client = storage_api.StorageClient()
      obj_ref = storage_util.ObjectReference.FromUrl(template_file_gcs_location)
      return storage_client.CopyFileToGCS(local_path, obj_ref)

  @staticmethod
  def BuildAndStoreFlexTemplateFile(template_file_gcs_location,
                                    image,
                                    template_metadata_json,
                                    sdk_language,
                                    print_only,
                                    template_args=None,
                                    image_repository_username_secret_id=None,
                                    image_repository_password_secret_id=None,
                                    image_repository_cert_path=None):
    """Builds container spec and stores it in the flex template file in GCS.

    Args:
      template_file_gcs_location: GCS location to store the template file.
      image: Path to the container image.
      template_metadata_json: Template metadata in json format.
      sdk_language: SDK language of the flex template.
      print_only: Only prints the container spec and skips write to GCS.
      template_args: Default runtime parameters specified by template authors.
      image_repository_username_secret_id: Secret manager secret id for username
        to authenticate to private registry.
      image_repository_password_secret_id: Secret manager secret id for password
        to authenticate to private registry.
      image_repository_cert_path: The full URL to self-signed certificate of
        private registry in Cloud Storage.

    Returns:
      Container spec json if print_only is set. A success message with template
      file GCS path and container spec otherewise.
    """
    template_metadata = None
    if template_metadata_json:
      template_metadata = Templates._BuildTemplateMetadata(
          template_metadata_json)
    sdk_info = Templates._BuildSDKInfo(sdk_language)
    default_environment = None
    if template_args:
      user_labels_list = Templates.__ConvertDictArguments(
          template_args.additional_user_labels,
          Templates.FLEX_TEMPLATE_USER_LABELS_VALUE)
      ip_private = Templates.IP_CONFIGURATION_ENUM_VALUE.WORKER_IP_PRIVATE
      ip_configuration = (
          ip_private if template_args.disable_public_ips else None
      )
      enable_streaming_engine = (
          True if template_args.enable_streaming_engine else None
      )
      default_environment = Templates.FLEX_TEMPLATE_ENVIRONMENT(
          serviceAccountEmail=template_args.service_account_email,
          maxWorkers=template_args.max_workers,
          numWorkers=template_args.num_workers,
          network=template_args.network,
          subnetwork=template_args.subnetwork,
          machineType=template_args.worker_machine_type,
          tempLocation=template_args.temp_location
          if template_args.temp_location else template_args.staging_location,
          stagingLocation=template_args.staging_location,
          kmsKeyName=template_args.kms_key_name,
          ipConfiguration=ip_configuration,
          workerRegion=template_args.worker_region,
          workerZone=template_args.worker_zone,
          enableStreamingEngine=enable_streaming_engine,
          additionalExperiments=(template_args.additional_experiments if
                                 template_args.additional_experiments else []),
          additionalUserLabels=Templates.FLEX_TEMPLATE_USER_LABELS_VALUE(
              additionalProperties=user_labels_list)
          if user_labels_list else None)
    container_spec = Templates.CONTAINER_SPEC(
        image=image,
        metadata=template_metadata,
        sdkInfo=sdk_info,
        defaultEnvironment=default_environment,
        imageRepositoryUsernameSecretId=image_repository_username_secret_id,
        imageRepositoryPasswordSecretId=image_repository_password_secret_id,
        imageRepositoryCertPath=image_repository_cert_path)
    container_spec_json = encoding.MessageToJson(container_spec)
    container_spec_pretty_json = json.dumps(
        json.loads(container_spec_json),
        sort_keys=True,
        indent=4,
        separators=(',', ': '))
    if print_only:
      return container_spec_pretty_json
    try:
      Templates._StoreFlexTemplateFile(template_file_gcs_location,
                                       container_spec_pretty_json)
      log.status.Print(
          'Successfully saved container spec in flex template file.\n'
          'Template File GCS Location: {}\n'
          'Container Spec:\n\n'
          '{}'.format(template_file_gcs_location, container_spec_pretty_json))
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def _AddPermissions(path, permissions):
    """Adds the given permissions to a file or directory.

    Args:
      path: The path to the file or directory.
      permissions: The permissions to add.

    Raises:
      OSError: If the chmod fails.
    """
    permissions = (
        os.stat(path).st_mode & Templates.ALL_PERMISSIONS_MASK
    ) | permissions
    os.chmod(path, permissions)

  @staticmethod
  def _ChmodRWorldReadable(top_dir_path):
    """Walks a dir to chmod itself and its contents with the configured access.

    Args:
      top_dir_path: The path to the top-level directory.

    Raises:
      OSError: If the chmod fails.
    """
    for dirpath, _, filenames in os.walk(top_dir_path):
      Templates._AddPermissions(dirpath, Templates.DIR_PERMISSIONS_MASK)
      for filename in filenames:
        Templates._AddPermissions(
            os.path.join(dirpath, filename), Templates.FILE_PERMISSIONS_MASK
        )

  @staticmethod
  def BuildAndStoreFlexTemplateImage(
      image_gcr_path,
      flex_template_base_image,
      jar_paths,
      py_paths,
      go_binary_path,
      env,
      sdk_language,
      gcs_log_dir,
      cloud_build_service_account,
  ):
    """Builds the flex template docker container image and stores it in GCR.

    Args:
      image_gcr_path: GCR location to store the flex template container image.
      flex_template_base_image: SDK version or base image to use.
      jar_paths: List of jar paths to pipelines and dependencies.
      py_paths: List of python paths to pipelines and dependencies.
      go_binary_path: Path to compiled Go pipeline binary.
      env: Dictionary of env variables to set in the container image.
      sdk_language: SDK language of the flex template.
      gcs_log_dir: Path to Google Cloud Storage directory to store build logs.
      cloud_build_service_account: Service account to be used by Cloud
        Build to build the image.

    Returns:
      True if container is built and store successfully.

    Raises:
      ValueError: If the parameters values are invalid.
    """
    Templates.__ValidateFlexTemplateEnv(env, sdk_language)
    with files.TemporaryDirectory() as temp_dir:
      log.status.Print('Copying files to a temp directory {}'.format(temp_dir))

      pipeline_files = []
      paths = jar_paths
      if py_paths:
        paths = py_paths
      elif go_binary_path:
        paths = [go_binary_path]

      for path in paths:
        absl_path = os.path.abspath(path)
        if os.path.isfile(absl_path):
          copy_file = shutil.copy2(absl_path, temp_dir)
          # Add the configured access to support non-root container execution.
          try:
            Templates._AddPermissions(
                copy_file, Templates.FILE_PERMISSIONS_MASK
            )
          except OSError:
            log.warning(
                'Could not adjust permissions for copied file {}'.format(
                    copy_file
                )
            )
        else:
          copy_dir = shutil.copytree(
              absl_path,
              os.path.join(temp_dir, os.path.basename(absl_path)),
          )
          # Add the configured access to support non-root container execution.
          try:
            Templates._ChmodRWorldReadable(copy_dir)
          except OSError:
            log.warning(
                'Could not adjust permissions for copied directory {}'.format(
                    copy_dir
                )
            )
        pipeline_files.append(os.path.split(absl_path)[1])

      log.status.Print(
          'Generating dockerfile to build the flex template container image...')
      dockerfile_contents = Templates.BuildDockerfile(flex_template_base_image,
                                                      pipeline_files, env,
                                                      sdk_language)

      dockerfile_path = os.path.join(temp_dir, 'Dockerfile')
      files.WriteFileContents(dockerfile_path, dockerfile_contents)
      log.status.Print(
          'Generated Dockerfile. Contents: {}'.format(dockerfile_contents))

      cloudbuild_messages = cloudbuild_util.GetMessagesModule()
      build_config = submit_util.CreateBuildConfig(
          tag=image_gcr_path,
          no_cache=False,
          messages=cloudbuild_messages,
          substitutions=None,
          arg_config='cloudbuild.yaml',
          is_specified_source=True,
          no_source=False,
          source=temp_dir,
          gcs_source_staging_dir=None,
          ignore_file=None,
          arg_gcs_log_dir=gcs_log_dir,
          arg_machine_type=None,
          arg_disk_size=None,
          arg_worker_pool=None,
          arg_dir=None,
          arg_revision=None,
          arg_git_source_dir=None,
          arg_git_source_revision=None,
          arg_service_account=(
              cloud_build_service_account
              if cloud_build_service_account
              else None
          ),
          buildpack=None,
      )
      log.status.Print('Pushing flex template container image to GCR...')

      submit_util.Build(cloudbuild_messages, False, build_config)
      return True

  @staticmethod
  def CreateJobFromFlexTemplate(template_args=None):
    """Calls the create job from flex template APIs.

    Args:
      template_args: Arguments for create template.

    Returns:
      (Job)
    """

    params_list = Templates.__ConvertDictArguments(
        template_args.parameters, Templates.FLEX_TEMPLATE_PARAMETERS_VALUE)
    transform_mapping_list = Templates.__ConvertDictArguments(
        template_args.transform_name_mappings,
        Templates.FLEX_TEMPLATE_TRANSFORM_NAME_MAPPING_VALUE)
    transform_mappings = None
    streaming_update = None
    if template_args.streaming_update:
      streaming_update = template_args.streaming_update
      if transform_mapping_list:
        transform_mappings = (
            Templates.FLEX_TEMPLATE_TRANSFORM_NAME_MAPPING_VALUE(
                additionalProperties=transform_mapping_list
            )
        )
    launch_options_list = []
    if template_args.launcher_vm_timeout_secs:
      launch_options_list.append(
          Templates.FLEX_TEMPLATE_LAUNCH_OPTIONS_VALUE.AdditionalProperty(
              key='ft_launch_timeout_secs',
              value=str(template_args.launcher_vm_timeout_secs),
          )
      )
    user_labels_list = Templates.__ConvertDictArguments(
        template_args.additional_user_labels,
        Templates.FLEX_TEMPLATE_USER_LABELS_VALUE)

    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = template_args.region_id or DATAFLOW_API_DEFAULT_REGION

    ip_private = Templates.IP_CONFIGURATION_ENUM_VALUE.WORKER_IP_PRIVATE
    ip_configuration = ip_private if template_args.disable_public_ips else None

    flexrs_goal = None
    if template_args.flexrs_goal:
      if template_args.flexrs_goal == 'SPEED_OPTIMIZED':
        flexrs_goal = Templates.FLEXRS_GOAL_ENUM_VALUE.FLEXRS_SPEED_OPTIMIZED
      elif template_args.flexrs_goal == 'COST_OPTIMIZED':
        flexrs_goal = Templates.FLEXRS_GOAL_ENUM_VALUE.FLEXRS_COST_OPTIMIZED

    body = Templates.LAUNCH_FLEX_TEMPLATE_REQUEST(
        launchParameter=Templates.FLEX_TEMPLATE_PARAMETER(
            jobName=template_args.job_name,
            containerSpecGcsPath=template_args.gcs_location,
            launchOptions=Templates.FLEX_TEMPLATE_LAUNCH_OPTIONS_VALUE(
                additionalProperties=launch_options_list
            ) if launch_options_list else None,
            environment=Templates.FLEX_TEMPLATE_ENVIRONMENT(
                serviceAccountEmail=template_args.service_account_email,
                maxWorkers=template_args.max_workers,
                numWorkers=template_args.num_workers,
                network=template_args.network,
                subnetwork=template_args.subnetwork,
                machineType=template_args.worker_machine_type,
                launcherMachineType=template_args.launcher_machine_type,
                tempLocation=template_args.temp_location if template_args
                .temp_location else template_args.staging_location,
                stagingLocation=template_args.staging_location,
                kmsKeyName=template_args.kms_key_name,
                ipConfiguration=ip_configuration,
                workerRegion=template_args.worker_region,
                workerZone=template_args.worker_zone,
                enableStreamingEngine=template_args.enable_streaming_engine,
                flexrsGoal=flexrs_goal,
                additionalExperiments=(
                    template_args.additional_experiments if template_args
                    .additional_experiments else []),
                additionalUserLabels=Templates.FLEX_TEMPLATE_USER_LABELS_VALUE(
                    additionalProperties=user_labels_list
                ) if user_labels_list else None,
                additionalPipelineOptions=(
                    template_args.additional_pipeline_options
                    if template_args.additional_pipeline_options
                    else []
                )),
            update=streaming_update,
            transformNameMappings=transform_mappings,
            parameters=Templates.FLEX_TEMPLATE_PARAMETERS_VALUE(
                additionalProperties=params_list) if params_list else None))
    request = GetMessagesModule(
    ).DataflowProjectsLocationsFlexTemplatesLaunchRequest(
        projectId=template_args.project_id or GetProject(),
        location=region_id,
        launchFlexTemplateRequest=body)
    try:
      return Templates.GetFlexTemplateService().Launch(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)


class Messages:
  """The Messages set of Dataflow API functions."""

  LIST_REQUEST = GetMessagesModule(
  ).DataflowProjectsLocationsJobsMessagesListRequest

  @staticmethod
  def GetService():
    return GetClientInstance().projects_locations_jobs_messages

  @staticmethod
  def List(job_id,
           project_id=None,
           region_id=None,
           minimum_importance=None,
           start_time=None,
           end_time=None,
           page_size=None,
           page_token=None):
    """Calls the Dataflow Metrics.Get method.

    Args:
      job_id: The job to get messages about.
      project_id: The project which owns the job.
      region_id: The regional endpoint of the job.
      minimum_importance: Filter to only get messages with importance >= level
      start_time: If specified, return only messages with timestamps >=
        start_time. The default is the job creation time (i.e. beginning of
        messages).
      end_time: Return only messages with timestamps < end_time. The default is
        now (i.e. return up to the latest messages available).
      page_size: If specified, determines the maximum number of messages to
        return.  If unspecified, the service may choose an appropriate default,
        or may return an arbitrarily large number of results.
      page_token: If supplied, this should be the value of next_page_token
        returned by an earlier call. This will cause the next page of results to
        be returned.

    Returns:
      (ListJobMessagesResponse)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    request = GetMessagesModule(
    ).DataflowProjectsLocationsJobsMessagesListRequest(
        jobId=job_id,
        location=region_id,
        projectId=project_id,
        startTime=start_time,
        endTime=end_time,
        minimumImportance=minimum_importance,
        pageSize=page_size,
        pageToken=page_token)
    try:
      return Messages.GetService().List(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)


class Snapshots:
  """Cloud Dataflow snapshots api."""

  @staticmethod
  def GetService():
    return GetClientInstance().projects_locations_snapshots

  @staticmethod
  def Delete(snapshot_id=None, project_id=None, region_id=None):
    """Calls the Dataflow Snapshots.Delete method.

    Args:
      snapshot_id: The id of the snapshot to delete.
      project_id: The project that owns the snapshot.
      region_id: The regional endpoint of the snapshot.

    Returns:
      (DeleteSnapshotResponse)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    request = GetMessagesModule(
    ).DataflowProjectsLocationsSnapshotsDeleteRequest(
        snapshotId=snapshot_id, location=region_id, projectId=project_id)
    try:
      return Snapshots.GetService().Delete(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def Get(snapshot_id=None, project_id=None, region_id=None):
    """Calls the Dataflow Snapshots.Get method.

    Args:
      snapshot_id: The id of the snapshot to get.
      project_id: The project that owns the snapshot.
      region_id: The regional endpoint of the snapshot.

    Returns:
      (GetSnapshotResponse)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    request = GetMessagesModule().DataflowProjectsLocationsSnapshotsGetRequest(
        snapshotId=snapshot_id, location=region_id, projectId=project_id)
    try:
      return Snapshots.GetService().Get(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)

  @staticmethod
  def List(job_id=None, project_id=None, region_id=None):
    """Calls the Dataflow Snapshots.List method.

    Args:
      job_id: If specified, only snapshots associated with the job will be
        returned.
      project_id: The project that owns the snapshot.
      region_id: The regional endpoint of the snapshot.

    Returns:
      (ListSnapshotsResponse)
    """
    project_id = project_id or GetProject()
    # TODO(b/139889563): Remove default when args region is changed to required
    region_id = region_id or DATAFLOW_API_DEFAULT_REGION
    request = GetMessagesModule().DataflowProjectsLocationsSnapshotsListRequest(
        jobId=job_id, location=region_id, projectId=project_id)
    try:
      return Snapshots.GetService().List(request)
    except apitools_exceptions.HttpError as error:
      raise exceptions.HttpException(error)
