<s> import logging <EOL> import time <EOL> from luigi_gcloud . gcore import _GCloudTask <EOL> from luigi_gcloud . storage import GCSFileSystem <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> default_pig_udf = None <EOL> default_spark_udf = None <EOL> class _DataProcJob : <EOL> def __init__ ( self , dataproc_api , project_id , job ) : <EOL> self . dataproc_api = dataproc_api <EOL> self . project_id = project_id <EOL> self . job = dataproc_api . projects ( ) . jobs ( ) . submit ( <EOL> projectId = self . project_id , <EOL> body = job ) . execute ( ) <EOL> self . job_id = self . job [ '<STR_LIT>' ] [ '<STR_LIT>' ] <EOL> logger . info ( '<STR_LIT>' , self . job_id , str ( self . job [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] ) ) <EOL> def wait_for_done ( self ) : <EOL> while True : <EOL> self . job = self . dataproc_api . projects ( ) . jobs ( ) . get ( <EOL> projectId = self . project_id , <EOL> jobId = self . job_id ) . execute ( ) <EOL> if '<STR_LIT>' == self . job [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] : <EOL> print ( str ( self . job ) ) <EOL> logger . error ( '<STR_LIT>' , self . job_id ) <EOL> logger . error ( self . job [ '<STR_LIT:status>' ] [ '<STR_LIT>' ] ) <EOL> logger . debug ( str ( self . job ) ) <EOL> return False <EOL> if '<STR_LIT>' == self . job [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] : <EOL> print ( str ( self . job ) ) <EOL> logger . warning ( '<STR_LIT>' , self . job_id ) <EOL> logger . warning ( self . job [ '<STR_LIT:status>' ] [ '<STR_LIT>' ] ) <EOL> logger . debug ( str ( self . job ) ) <EOL> return False <EOL> if '<STR_LIT>' == self . job [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] : <EOL> return True <EOL> logger . debug ( '<STR_LIT>' , self . job_id , str ( self . job [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] ) ) <EOL> time . sleep ( <NUM_LIT:5> ) <EOL> def raise_error ( self , message = None ) : <EOL> if '<STR_LIT>' == self . job [ '<STR_LIT:status>' ] [ '<STR_LIT:state>' ] : <EOL> if message is None : <EOL> message = "<STR_LIT>" <EOL> raise Exception ( message + "<STR_LIT>" + str ( self . job [ '<STR_LIT:status>' ] [ '<STR_LIT>' ] ) ) <EOL> def get ( self ) : <EOL> return self . job <EOL> class DataProcPigTask ( _GCloudTask ) : <EOL> service_name = "<STR_LIT>" <EOL> def query_file ( self ) : <EOL> return None <EOL> def query_uri ( self ) : <EOL> return self . client . project_staging ( ) + self . resolved_name ( ) + "<STR_LIT>" <EOL> def lib_uris ( self ) : <EOL> return [ ] <EOL> def _jar_file_uris ( self ) : <EOL> global default_pig_udf <EOL> artifacts = [ ] <EOL> uris = self . lib_uris ( ) <EOL> artifacts . extend ( uris ) <EOL> if default_pig_udf is not None : <EOL> if not default_pig_udf . get ( '<STR_LIT>' ) and len ( uris ) == <NUM_LIT:0> : <EOL> artifacts . extend ( default_pig_udf . get ( '<STR_LIT>' ) ) <EOL> elif default_pig_udf . get ( '<STR_LIT>' ) : <EOL> artifacts . extend ( default_pig_udf . get ( '<STR_LIT>' ) ) <EOL> return artifacts <EOL> def run ( self ) : <EOL> self . start_job ( ) <EOL> http = self . client . http_authorized ( ) <EOL> dataproc_api = self . client . dataproc_api ( http ) <EOL> name = self . job_name ( ) <EOL> if self . query_file ( ) is not None : <EOL> fs = GCSFileSystem ( ) <EOL> fs . put ( self . query_file ( ) , <EOL> self . query_uri ( ) ) <EOL> job = { <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . client . project_id ( ) , <EOL> "<STR_LIT>" : name , <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . get_service_value ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT>" : self . query_uri ( ) , <EOL> "<STR_LIT>" : self . _jar_file_uris ( ) , <EOL> } <EOL> } <EOL> } <EOL> variables = self . variables ( ) <EOL> if variables is not None : <EOL> job [ "<STR_LIT>" ] [ "<STR_LIT>" ] [ "<STR_LIT>" ] = variables <EOL> print ( job ) <EOL> submitted = _DataProcJob ( dataproc_api , self . client . project_id ( ) , job ) <EOL> if not submitted . wait_for_done ( ) : <EOL> submitted . raise_error ( "<STR_LIT>" ) <EOL> class DataProcSparkTask ( _GCloudTask ) : <EOL> service_name = "<STR_LIT>" <EOL> def job_file ( self ) : <EOL> return None <EOL> def job_uri ( self ) : <EOL> return self . client . project_staging ( ) + self . resolved_name ( ) + "<STR_LIT>" <EOL> def lib_uris ( self ) : <EOL> return [ ] <EOL> def main ( self ) : <EOL> raise NotImplementedError <EOL> def args ( self ) : <EOL> return [ ] <EOL> def _jar_file_uris ( self ) : <EOL> global default_spark_udf <EOL> artifacts = [ self . job_uri ( ) ] <EOL> uris = self . lib_uris ( ) <EOL> artifacts . extend ( uris ) <EOL> if default_spark_udf is not None : <EOL> if not default_spark_udf . get ( '<STR_LIT>' ) and len ( uris ) == <NUM_LIT:0> : <EOL> artifacts . extend ( default_spark_udf . get ( '<STR_LIT>' ) ) <EOL> elif default_spark_udf . get ( '<STR_LIT>' ) : <EOL> artifacts . extend ( default_spark_udf . get ( '<STR_LIT>' ) ) <EOL> return artifacts <EOL> def run ( self ) : <EOL> self . start_job ( ) <EOL> http = self . client . http_authorized ( ) <EOL> dataproc_api = self . client . dataproc_api ( http ) <EOL> name = self . job_name ( ) <EOL> if self . job_file ( ) is not None : <EOL> logger . warning ( "<STR_LIT>" ) <EOL> fs = GCSFileSystem ( ) <EOL> fs . put ( self . get_service_value ( "<STR_LIT>" , "<STR_LIT:.>" ) + self . job_file ( ) , <EOL> self . job_uri ( ) ) <EOL> job = { <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . client . project_id ( ) , <EOL> "<STR_LIT>" : name , <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . get_service_value ( "<STR_LIT>" , "<STR_LIT>" ) <EOL> } , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : self . main ( ) , <EOL> "<STR_LIT>" : self . _jar_file_uris ( ) , <EOL> } <EOL> } <EOL> } <EOL> args = self . args ( ) <EOL> if args is not None : <EOL> job [ "<STR_LIT>" ] [ "<STR_LIT>" ] [ "<STR_LIT:args>" ] = args <EOL> print ( job ) <EOL> submitted = _DataProcJob ( dataproc_api , self . client . project_id ( ) , job ) <EOL> if not submitted . wait_for_done ( ) : <EOL> submitted . raise_error ( "<STR_LIT>" ) <EOL> def set_default_pig_udf ( uris , append = False ) : <EOL> global default_pig_udf <EOL> default_pig_udf = { <EOL> '<STR_LIT>' : append , <EOL> '<STR_LIT>' : uris <EOL> } <EOL> def set_default_spark_udf ( uris , append = False ) : <EOL> global default_spark_udf <EOL> default_spark_udf = { <EOL> '<STR_LIT>' : append , <EOL> '<STR_LIT>' : uris <EOL> } </s>
