
    ]                         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  G d	 d
ej                        Zy)zBase class for the Trino job.    )absolute_import)division)unicode_literals)encoding)arg_parsers)base)utilc                   @    e Zd ZdZed        Zed        Zed        Zy)	TrinoBasez Submit a Trino job to a cluster.c           	         | j                  d      }|j                  dddddg d	       |j                  d
dd       | j                  dt        j                         dd       | j                  dt        j
                         | j                  dt        j                         dd       | j                  ddd       | j                  dd       | j                  dt        j                         dd       y) z@Parses command line arguments specific to submitting Trino jobs.T)requiredz	--executez-eQUERYqueriesappendzA Trino query to execute.)metavardestactiondefaulthelpz--filez-fz8HCFS URI of file containing the Trino script to execute.)r   z--propertieszPARAM=VALUEz:A list of key value pairs to set Trino session properties.)typer   r   z--properties-filez--driver-log-levelszPACKAGE=LEVELzpA list of package-to-log4j log level pairs to configure driver logging. For example: root=FATAL,com.example=INFOz--continue-on-failure
store_truez%Whether to continue if a query fails.)r   r   z--query-output-formatzZThe query output display format. See the Trino documentation for supported output formats.z--client-tags
CLIENT_TAGz4A list of Trino client tags to attach to this query.N)add_mutually_exclusive_groupadd_argumentr   ArgDictjob_utilPROPERTIES_FILE_HELP_TEXTArgList)parserdrivers     5lib/googlecloudsdk/command_lib/dataproc/jobs/trino.pyArgszTrinoBase.Args   s>    00$0?F
(  * G  I   "I	  K
 //  1   "B	  D 4  6 .  0   "C	  E    c                     d| j                   iS )Nfile)r%   )argss    r!   GetFilesByTypezTrinoBase.GetFilesByTypeK   s    DIIr#   c                    | j                  |j                  |d   |      }|j                  r!| j                  |j                        |_        |j
                  r|j
                  |_        |j                  r|j                  |_        t        j                  |j                  |j                        }|r1t        j                  || j                   j                  d      |_        ||_        y)z/Populates the trinoJob member of the given job.r%   )continueOnFailurequeryFileUriloggingConfig)r   T)
sort_itemsN)TrinoJobcontinue_on_failurer   	QueryList	queryListquery_output_formatoutputFormatclient_tags
clientTagsr   BuildJobProperties
propertiesproperties_filer   DictToAdditionalPropertyMessagePropertiesValuetrinoJob)messagesjobfiles_by_typelogging_configr&   	trino_jobjob_propertiess          r!   ConfigureJobzTrinoBase.ConfigureJobO   s     !!22"6*$ " &I
 ||$..t||.Di#77i!--i00--/N%EE
(++;;Ni CLr#   N)__name__
__module____qualname____doc__staticmethodr"   r'   rA    r#   r!   r   r      sB    ()E )EV    r#   r   N)rE   
__future__r   r   r   apitools.base.pyr   googlecloudsdk.callioper   (googlecloudsdk.command_lib.dataproc.jobsr   job_baser	   r   JobBaser   rG   r#   r!   <module>rN      s3    $ &  ' % / E EJ   Jr#   