
    -                         d Z ddlmZ ddlmZ ddlmZ ddlmZ ddlmZ ddl	m
Z
 ddlmZ dd	lmZ dd
lmZ e
j                    G d de
j"                               Zy)z2Set managed cluster for workflow template command.    )absolute_import)division)unicode_literals)compute_helpers)dataproc)base)clusters)flags)labels_utilc                   ^    e Zd ZdZddiZed        Zed        Zed        Zed        Z	d Z
y	)
SetManagedClusterz0Set a managed cluster for the workflow template.EXAMPLESz
To update managed cluster in a workflow template, run:

  $ {command} my_template --region=us-central1 --no-address --num-workers=10 --worker-machine-type=custom-6-23040

c                    t        j                  | j                               }|j                  dd       t	        j
                  ||| j                         | j                         | j                         d       t        j                  |d|j                         | j                         rt	        j                  |       y y )Nz--cluster-namezt          The name of the managed dataproc cluster.
          If unspecified, the workflow template ID will be used.)helpF)include_deprecatedinclude_gke_platform_argszset managed cluster)dpDataprocReleaseTrackadd_argumentr	   ArgsForClusterRefBetaAlphar
   AddTemplateResourceArgapi_versionBetaArgsForClusterRef)clsparserr   s      >lib/surface/dataproc/workflow_templates/set_managed_cluster.pyArgszSetManagedCluster.Args+   s    {{3++-.H
D  E
 
		88:"') 
  )>!)!5!57
xxz$$V,     c                 X    | j                         t        j                   j                  k7  S N)r   r   GAr   s    r   r   zSetManagedCluster.Beta?   s"    !2!2!5!555r!   c                 X    | j                         t        j                   j                  k(  S r#   )r   r   ALPHAr%   s    r   r   zSetManagedCluster.AlphaC   s"    !2!2!8!888r!   c                     | j                         rt        j                  j                  S t        j                  j                  S r#   )r   r   r   BETAr$   r%   s    r   GetComputeReleaseTrackz(SetManagedCluster.GetComputeReleaseTrackG   s/    
xxz###r!   c           
      D   t        j                  | j                               }|j                  j                  j                         }|j                  ||j                        }|j                  r|j                  }n|j                  }t        j                  | j                         ||j                        }t        j                  |||j                   || j#                         | j%                         | j#                               }t'        j(                  ||j*                  j,                  j.                        }|j*                  j-                  |||      }	|j*                  j1                  |	      |_        |j4                  j6                  j9                  |      }
|
S )N)r   )clusterNameconfiglabels)managedCluster)r   r   r   CONCEPTStemplateParseGetRegionsWorkflowTemplateversioncluster_nameworkflowTemplatesIdr   GetComputeResourcesr*   	regionsIdr	   GetClusterConfig
projectsIdr   r   r   ParseCreateArgsmessagesManagedClusterLabelsValueWorkflowTemplatePlacement	placementclient"projects_regions_workflowTemplatesUpdate)selfargsr   template_refworkflow_templater5   compute_resourcescluster_configr.   managed_clusterresponses              r   RunzSetManagedCluster.RunM   sj   {{4,,./H==))//1L ;;dll$ &&l!55l';;##%|\5K5KM ..		

99;(N ((h..::<F ''66  7 HO #+"3"3"M"M& #N #( AAHHHOr!   N)__name__
__module____qualname____doc__detailed_helpclassmethodr    r   r   r*   rL    r!   r   r   r      sk    8  - - -& 6 6 9 9    
$r!   r   N)rP   
__future__r   r   r   googlecloudsdk.api_lib.dataprocr   r   r   googlecloudsdk.callioper   #googlecloudsdk.command_lib.dataprocr	   r
   $googlecloudsdk.command_lib.util.argsr   DefaultUniverseOnlyUpdateCommandr   rS   r!   r   <module>r[      sO    9 &  ' ; : ( 8 5 < S** S Sr!   