<s> import logging <EOL> import re <EOL> from airflow . hooks import HiveCliHook <EOL> from airflow . models import BaseOperator <EOL> from airflow . utils . decorators import apply_defaults <EOL> class HiveOperator ( BaseOperator ) : <EOL> """<STR_LIT>""" <EOL> template_fields = ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> template_ext = ( '<STR_LIT>' , '<STR_LIT>' , ) <EOL> ui_color = '<STR_LIT>' <EOL> @ apply_defaults <EOL> def __init__ ( <EOL> self , hql , <EOL> hive_cli_conn_id = '<STR_LIT>' , <EOL> schema = '<STR_LIT:default>' , <EOL> hiveconf_jinja_translate = False , <EOL> script_begin_tag = None , <EOL> run_as_owner = False , <EOL> * args , ** kwargs ) : <EOL> super ( HiveOperator , self ) . __init__ ( * args , ** kwargs ) <EOL> self . hiveconf_jinja_translate = hiveconf_jinja_translate <EOL> self . hql = hql <EOL> self . schema = schema <EOL> self . hive_cli_conn_id = hive_cli_conn_id <EOL> self . script_begin_tag = script_begin_tag <EOL> self . run_as = None <EOL> if run_as_owner : <EOL> self . run_as = self . dag . owner <EOL> def get_hook ( self ) : <EOL> return HiveCliHook ( hive_cli_conn_id = self . hive_cli_conn_id , run_as = self . run_as ) <EOL> def prepare_template ( self ) : <EOL> if self . hiveconf_jinja_translate : <EOL> self . hql = re . sub ( <EOL> "<STR_LIT>" , "<STR_LIT>" , self . hql ) <EOL> if self . script_begin_tag and self . script_begin_tag in self . hql : <EOL> self . hql = "<STR_LIT:\n>" . join ( self . hql . split ( self . script_begin_tag ) [ <NUM_LIT:1> : ] ) <EOL> def execute ( self , context ) : <EOL> logging . info ( '<STR_LIT>' + self . hql ) <EOL> self . hook = self . get_hook ( ) <EOL> self . hook . run_cli ( hql = self . hql , schema = self . schema ) <EOL> def dry_run ( self ) : <EOL> self . hook = self . get_hook ( ) <EOL> self . hook . test_hql ( hql = self . hql ) <EOL> def on_kill ( self ) : <EOL> self . hook . kill ( ) </s>
