<s> import os <EOL> import time <EOL> import datetime <EOL> import logging <EOL> from airflow . exceptions import AirflowException <EOL> from airflow . hooks . base_hook import BaseHook <EOL> from airflow import configuration <EOL> from qds_sdk . qubole import Qubole <EOL> from qds_sdk . commands import Command , HiveCommand , PrestoCommand , HadoopCommand , PigCommand , ShellCommand , SparkCommand , DbTapQueryCommand , DbExportCommand , DbImportCommand <EOL> COMMAND_CLASSES = { <EOL> "<STR_LIT>" : HiveCommand , <EOL> "<STR_LIT>" : PrestoCommand , <EOL> "<STR_LIT>" : HadoopCommand , <EOL> "<STR_LIT>" : ShellCommand , <EOL> "<STR_LIT>" : PigCommand , <EOL> "<STR_LIT>" : SparkCommand , <EOL> "<STR_LIT>" : DbTapQueryCommand , <EOL> "<STR_LIT>" : DbExportCommand , <EOL> "<STR_LIT>" : DbImportCommand <EOL> } <EOL> HYPHEN_ARGS = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> POSITIONAL_ARGS = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> COMMAND_ARGS = { <EOL> "<STR_LIT>" : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' , '<STR_LIT>' , '<STR_LIT>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' ] , <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:name>' ] <EOL> } <EOL> class QuboleHook ( BaseHook ) : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> conn = self . get_connection ( kwargs [ '<STR_LIT>' ] ) <EOL> Qubole . configure ( api_token = conn . password , api_url = conn . host ) <EOL> self . task_id = kwargs [ '<STR_LIT>' ] <EOL> self . dag_id = kwargs [ '<STR_LIT>' ] . dag_id <EOL> self . kwargs = kwargs <EOL> self . args = self . create_cmd_args ( ) <EOL> self . cls = COMMAND_CLASSES [ self . kwargs [ '<STR_LIT>' ] ] <EOL> self . cmd = None <EOL> def execute ( self , context ) : <EOL> args = self . cls . parse ( self . args ) <EOL> self . cmd = self . cls . create ( ** args ) <EOL> context [ '<STR_LIT>' ] . xcom_push ( key = '<STR_LIT>' , value = self . cmd . id ) <EOL> logging . info ( "<STR_LIT>" . format ( str ( self . cmd . id ) , self . cmd . status ) ) <EOL> while not Command . is_done ( self . cmd . status ) : <EOL> time . sleep ( Qubole . poll_interval ) <EOL> self . cmd = self . cls . find ( self . cmd . id ) <EOL> logging . info ( "<STR_LIT>" . format ( str ( self . cmd . id ) , self . cmd . status ) ) <EOL> if self . kwargs . has_key ( '<STR_LIT>' ) and self . kwargs [ '<STR_LIT>' ] == True : <EOL> logging . info ( "<STR_LIT>" . format ( str ( self . cmd . id ) , self . cmd . get_log ( ) ) ) <EOL> if self . cmd . status != '<STR_LIT>' : <EOL> raise AirflowException ( '<STR_LIT>' . format ( self . cmd . id , self . cmd . status ) ) <EOL> def kill ( self , ti ) : <EOL> """<STR_LIT>""" <EOL> if self . cmd is None : <EOL> cmd_id = ti . xcom_pull ( key = "<STR_LIT>" , task_ids = self . task_id ) <EOL> self . cmd = self . cls . find ( cmd_id ) <EOL> if self . cls and self . cmd : <EOL> logging . info ( '<STR_LIT>' . format ( self . cmd . id ) ) <EOL> self . cmd . cancel ( ) <EOL> def get_results ( self , ti = None , fp = None , inline = True , delim = None , fetch = True ) : <EOL> """<STR_LIT>""" <EOL> if fp is None : <EOL> iso = datetime . datetime . utcnow ( ) . isoformat ( ) <EOL> logpath = os . path . expanduser ( configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> resultpath = logpath + '<STR_LIT:/>' + self . dag_id + '<STR_LIT:/>' + self . task_id + '<STR_LIT>' <EOL> configuration . mkdir_p ( resultpath ) <EOL> fp = open ( resultpath + '<STR_LIT:/>' + iso , '<STR_LIT:wb>' ) <EOL> if self . cmd is None : <EOL> cmd_id = ti . xcom_pull ( key = "<STR_LIT>" , task_ids = self . task_id ) <EOL> self . cmd = self . cls . find ( cmd_id ) <EOL> self . cmd . get_results ( fp , inline , delim , fetch ) <EOL> fp . flush ( ) <EOL> fp . close ( ) <EOL> return fp . name <EOL> def get_log ( self , ti ) : <EOL> """<STR_LIT>""" <EOL> if self . cmd is None : <EOL> cmd_id = ti . xcom_pull ( key = "<STR_LIT>" , task_ids = self . task_id ) <EOL> Command . get_log_id ( self . cls , cmd_id ) <EOL> def get_jobs_id ( self , ti ) : <EOL> """<STR_LIT>""" <EOL> if self . cmd is None : <EOL> cmd_id = ti . xcom_pull ( key = "<STR_LIT>" , task_ids = self . task_id ) <EOL> Command . get_jobs_id ( self . cls , cmd_id ) <EOL> def create_cmd_args ( self ) : <EOL> args = [ ] <EOL> cmd_type = self . kwargs [ '<STR_LIT>' ] <EOL> inplace_args = None <EOL> for k , v in self . kwargs . items ( ) : <EOL> if k in COMMAND_ARGS [ cmd_type ] : <EOL> if k in HYPHEN_ARGS : <EOL> args . append ( "<STR_LIT>" . format ( k . replace ( '<STR_LIT:_>' , '<STR_LIT:->' ) , v ) ) <EOL> elif k in POSITIONAL_ARGS : <EOL> inplace_args = v <EOL> else : <EOL> args . append ( "<STR_LIT>" . format ( k , v ) ) <EOL> if k == '<STR_LIT>' and v is True : <EOL> args . append ( "<STR_LIT>" ) <EOL> if inplace_args is not None : <EOL> if cmd_type == '<STR_LIT>' : <EOL> args += inplace_args . split ( '<STR_LIT:U+0020>' , <NUM_LIT:1> ) <EOL> else : <EOL> args += inplace_args . split ( '<STR_LIT:U+0020>' ) <EOL> return args </s>
