<s> from __future__ import print_function <EOL> from builtins import zip <EOL> from past . builtins import basestring <EOL> import unicodecsv as csv <EOL> import logging <EOL> import re <EOL> import subprocess <EOL> from tempfile import NamedTemporaryFile <EOL> from airflow . exceptions import AirflowException <EOL> from airflow . hooks . base_hook import BaseHook <EOL> from airflow . utils . file import TemporaryDirectory <EOL> from airflow import configuration <EOL> import airflow . security . utils as utils <EOL> class HiveCliHook ( BaseHook ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( <EOL> self , <EOL> hive_cli_conn_id = "<STR_LIT>" , <EOL> run_as = None ) : <EOL> conn = self . get_connection ( hive_cli_conn_id ) <EOL> self . hive_cli_params = conn . extra_dejson . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . use_beeline = conn . extra_dejson . get ( '<STR_LIT>' , False ) <EOL> self . auth = conn . extra_dejson . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . conn = conn <EOL> self . run_as = run_as <EOL> def run_cli ( self , hql , schema = None , verbose = True ) : <EOL> """<STR_LIT>""" <EOL> conn = self . conn <EOL> schema = schema or conn . schema <EOL> if schema : <EOL> hql = "<STR_LIT>" . format ( ** locals ( ) ) <EOL> with TemporaryDirectory ( prefix = '<STR_LIT>' ) as tmp_dir : <EOL> with NamedTemporaryFile ( dir = tmp_dir ) as f : <EOL> f . write ( hql . encode ( '<STR_LIT>' ) ) <EOL> f . flush ( ) <EOL> fname = f . name <EOL> hive_bin = '<STR_LIT>' <EOL> cmd_extra = [ ] <EOL> if self . use_beeline : <EOL> hive_bin = '<STR_LIT>' <EOL> jdbc_url = "<STR_LIT>" <EOL> if configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) == '<STR_LIT>' : <EOL> template = conn . extra_dejson . get ( <EOL> '<STR_LIT>' , "<STR_LIT>" ) <EOL> if "<STR_LIT>" in template : <EOL> template = utils . replace_hostname_pattern ( <EOL> utils . get_components ( template ) ) <EOL> proxy_user = "<STR_LIT>" <EOL> if conn . extra_dejson . get ( '<STR_LIT>' ) == "<STR_LIT>" and conn . login : <EOL> proxy_user = "<STR_LIT>" . format ( conn . login ) <EOL> elif conn . extra_dejson . get ( '<STR_LIT>' ) == "<STR_LIT>" and self . run_as : <EOL> proxy_user = "<STR_LIT>" . format ( self . run_as ) <EOL> jdbc_url += "<STR_LIT>" <EOL> elif self . auth : <EOL> jdbc_url += "<STR_LIT>" + self . auth <EOL> jdbc_url = jdbc_url . format ( ** locals ( ) ) <EOL> cmd_extra += [ '<STR_LIT>' , jdbc_url ] <EOL> if conn . login : <EOL> cmd_extra += [ '<STR_LIT>' , conn . login ] <EOL> if conn . password : <EOL> cmd_extra += [ '<STR_LIT>' , conn . password ] <EOL> hive_cmd = [ hive_bin , '<STR_LIT>' , fname ] + cmd_extra <EOL> if self . hive_cli_params : <EOL> hive_params_list = self . hive_cli_params . split ( ) <EOL> hive_cmd . extend ( hive_params_list ) <EOL> if verbose : <EOL> logging . info ( "<STR_LIT:U+0020>" . join ( hive_cmd ) ) <EOL> sp = subprocess . Popen ( <EOL> hive_cmd , <EOL> stdout = subprocess . PIPE , <EOL> stderr = subprocess . STDOUT , <EOL> cwd = tmp_dir ) <EOL> self . sp = sp <EOL> stdout = '<STR_LIT>' <EOL> while True : <EOL> line = sp . stdout . readline ( ) <EOL> if not line : <EOL> break <EOL> stdout += line . decode ( '<STR_LIT>' ) <EOL> if verbose : <EOL> logging . info ( line . decode ( '<STR_LIT>' ) . strip ( ) ) <EOL> sp . wait ( ) <EOL> if sp . returncode : <EOL> raise AirflowException ( stdout ) <EOL> return stdout <EOL> def test_hql ( self , hql ) : <EOL> """<STR_LIT>""" <EOL> create , insert , other = [ ] , [ ] , [ ] <EOL> for query in hql . split ( '<STR_LIT:;>' ) : <EOL> query_original = query <EOL> query = query . lower ( ) . strip ( ) <EOL> if query . startswith ( '<STR_LIT>' ) : <EOL> create . append ( query_original ) <EOL> elif query . startswith ( ( '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ) ) : <EOL> other . append ( query_original ) <EOL> elif query . startswith ( '<STR_LIT>' ) : <EOL> insert . append ( query_original ) <EOL> other = '<STR_LIT:;>' . join ( other ) <EOL> for query_set in [ create , insert ] : <EOL> for query in query_set : <EOL> query_preview = '<STR_LIT:U+0020>' . join ( query . split ( ) ) [ : <NUM_LIT:50> ] <EOL> logging . info ( "<STR_LIT>" . format ( query_preview ) ) <EOL> if query_set == insert : <EOL> query = other + '<STR_LIT>' + query <EOL> else : <EOL> query = '<STR_LIT>' + query <EOL> try : <EOL> self . run_cli ( query , verbose = False ) <EOL> except AirflowException as e : <EOL> message = e . args [ <NUM_LIT:0> ] . split ( '<STR_LIT:\n>' ) [ - <NUM_LIT:2> ] <EOL> logging . info ( message ) <EOL> error_loc = re . search ( '<STR_LIT>' , message ) <EOL> if error_loc and error_loc . group ( <NUM_LIT:1> ) . isdigit ( ) : <EOL> l = int ( error_loc . group ( <NUM_LIT:1> ) ) <EOL> begin = max ( l - <NUM_LIT:2> , <NUM_LIT:0> ) <EOL> end = min ( l + <NUM_LIT:3> , len ( query . split ( '<STR_LIT:\n>' ) ) ) <EOL> context = '<STR_LIT:\n>' . join ( query . split ( '<STR_LIT:\n>' ) [ begin : end ] ) <EOL> logging . info ( "<STR_LIT>" . format ( context ) ) <EOL> else : <EOL> logging . info ( "<STR_LIT>" ) <EOL> def load_file ( <EOL> self , <EOL> filepath , <EOL> table , <EOL> delimiter = "<STR_LIT:U+002C>" , <EOL> field_dict = None , <EOL> create = True , <EOL> overwrite = True , <EOL> partition = None , <EOL> recreate = False ) : <EOL> """<STR_LIT>""" <EOL> hql = '<STR_LIT>' <EOL> if recreate : <EOL> hql += "<STR_LIT>" <EOL> if create or recreate : <EOL> fields = "<STR_LIT>" . join ( <EOL> [ k + '<STR_LIT:U+0020>' + v for k , v in field_dict . items ( ) ] ) <EOL> hql += "<STR_LIT>" <EOL> if partition : <EOL> pfields = "<STR_LIT>" . join ( <EOL> [ p + "<STR_LIT>" for p in partition ] ) <EOL> hql += "<STR_LIT>" <EOL> hql += "<STR_LIT>" <EOL> hql += "<STR_LIT>" <EOL> hql += "<STR_LIT>" <EOL> hql = hql . format ( ** locals ( ) ) <EOL> logging . info ( hql ) <EOL> self . run_cli ( hql ) <EOL> hql = "<STR_LIT>" <EOL> if overwrite : <EOL> hql += "<STR_LIT>" <EOL> hql += "<STR_LIT>" <EOL> if partition : <EOL> pvals = "<STR_LIT:U+002CU+0020>" . join ( <EOL> [ "<STR_LIT>" . format ( k , v ) for k , v in partition . items ( ) ] ) <EOL> hql += "<STR_LIT>" <EOL> hql = hql . format ( ** locals ( ) ) <EOL> logging . info ( hql ) <EOL> self . run_cli ( hql ) <EOL> def kill ( self ) : <EOL> if hasattr ( self , '<STR_LIT>' ) : <EOL> if self . sp . poll ( ) is None : <EOL> print ( "<STR_LIT>" ) <EOL> self . sp . kill ( ) <EOL> class HiveMetastoreHook ( BaseHook ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , metastore_conn_id = '<STR_LIT>' ) : <EOL> self . metastore_conn = self . get_connection ( metastore_conn_id ) <EOL> self . metastore = self . get_metastore_client ( ) <EOL> def __getstate__ ( self ) : <EOL> d = dict ( self . __dict__ ) <EOL> del d [ '<STR_LIT>' ] <EOL> return d <EOL> def __setstate__ ( self , d ) : <EOL> self . __dict__ . update ( d ) <EOL> self . __dict__ [ '<STR_LIT>' ] = self . get_metastore_client ( ) <EOL> def get_metastore_client ( self ) : <EOL> """<STR_LIT>""" <EOL> from thrift . transport import TSocket , TTransport <EOL> from thrift . protocol import TBinaryProtocol <EOL> from hive_service import ThriftHive <EOL> ms = self . metastore_conn <EOL> auth_mechanism = ms . extra_dejson . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> if configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) == '<STR_LIT>' : <EOL> auth_mechanism = ms . extra_dejson . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> kerberos_service_name = ms . extra_dejson . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> socket = TSocket . TSocket ( ms . host , ms . port ) <EOL> if configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) == '<STR_LIT>' and auth_mechanism == '<STR_LIT>' : <EOL> try : <EOL> import saslwrapper as sasl <EOL> except ImportError : <EOL> import sasl <EOL> def sasl_factory ( ) : <EOL> sasl_client = sasl . Client ( ) <EOL> sasl_client . setAttr ( "<STR_LIT:host>" , ms . host ) <EOL> sasl_client ( "<STR_LIT>" , kerberos_service_name ) <EOL> sasl_client . init ( ) <EOL> from thrift_sasl import TSaslClientTransport <EOL> transport = TSaslClientTransport ( sasl_factory , "<STR_LIT>" , socket ) <EOL> else : <EOL> transport = TTransport . TBufferedTransport ( socket ) <EOL> protocol = TBinaryProtocol . TBinaryProtocol ( transport ) <EOL> return ThriftHive . Client ( protocol ) <EOL> def get_conn ( self ) : <EOL> return self . metastore <EOL> def check_for_partition ( self , schema , table , partition ) : <EOL> """<STR_LIT>""" <EOL> self . metastore . _oprot . trans . open ( ) <EOL> partitions = self . metastore . get_partitions_by_filter ( <EOL> schema , table , partition , <NUM_LIT:1> ) <EOL> self . metastore . _oprot . trans . close ( ) <EOL> if partitions : <EOL> return True <EOL> else : <EOL> return False <EOL> def get_table ( self , table_name , db = '<STR_LIT:default>' ) : <EOL> """<STR_LIT>""" <EOL> self . metastore . _oprot . trans . open ( ) <EOL> if db == '<STR_LIT:default>' and '<STR_LIT:.>' in table_name : <EOL> db , table_name = table_name . split ( '<STR_LIT:.>' ) [ : <NUM_LIT:2> ] <EOL> table = self . metastore . get_table ( dbname = db , tbl_name = table_name ) <EOL> self . metastore . _oprot . trans . close ( ) <EOL> return table <EOL> def get_tables ( self , db , pattern = '<STR_LIT:*>' ) : <EOL> """<STR_LIT>""" <EOL> self . metastore . _oprot . trans . open ( ) <EOL> tables = self . metastore . get_tables ( db_name = db , pattern = pattern ) <EOL> objs = self . metastore . get_table_objects_by_name ( db , tables ) <EOL> self . metastore . _oprot . trans . close ( ) <EOL> return objs <EOL> def get_databases ( self , pattern = '<STR_LIT:*>' ) : <EOL> """<STR_LIT>""" <EOL> self . metastore . _oprot . trans . open ( ) <EOL> dbs = self . metastore . get_databases ( pattern ) <EOL> self . metastore . _oprot . trans . close ( ) <EOL> return dbs <EOL> def get_partitions ( <EOL> self , schema , table_name , filter = None ) : <EOL> """<STR_LIT>""" <EOL> self . metastore . _oprot . trans . open ( ) <EOL> table = self . metastore . get_table ( dbname = schema , tbl_name = table_name ) <EOL> if len ( table . partitionKeys ) == <NUM_LIT:0> : <EOL> raise AirflowException ( "<STR_LIT>" ) <EOL> else : <EOL> if filter : <EOL> parts = self . metastore . get_partitions_by_filter ( <EOL> db_name = schema , tbl_name = table_name , <EOL> filter = filter , max_parts = <NUM_LIT> ) <EOL> else : <EOL> parts = self . metastore . get_partitions ( <EOL> db_name = schema , tbl_name = table_name , max_parts = <NUM_LIT> ) <EOL> self . metastore . _oprot . trans . close ( ) <EOL> pnames = [ p . name for p in table . partitionKeys ] <EOL> return [ dict ( zip ( pnames , p . values ) ) for p in parts ] <EOL> def max_partition ( self , schema , table_name , field = None , filter = None ) : <EOL> """<STR_LIT>""" <EOL> parts = self . get_partitions ( schema , table_name , filter ) <EOL> if not parts : <EOL> return None <EOL> elif len ( parts [ <NUM_LIT:0> ] ) == <NUM_LIT:1> : <EOL> field = list ( parts [ <NUM_LIT:0> ] . keys ( ) ) [ <NUM_LIT:0> ] <EOL> elif not field : <EOL> raise AirflowException ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> return max ( [ p [ field ] for p in parts ] ) <EOL> def table_exists ( self , table_name , db = '<STR_LIT:default>' ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> t = self . get_table ( table_name , db ) <EOL> return True <EOL> except Exception as e : <EOL> return False <EOL> class HiveServer2Hook ( BaseHook ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , hiveserver2_conn_id = '<STR_LIT>' ) : <EOL> self . hiveserver2_conn_id = hiveserver2_conn_id <EOL> def get_conn ( self ) : <EOL> db = self . get_connection ( self . hiveserver2_conn_id ) <EOL> auth_mechanism = db . extra_dejson . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> kerberos_service_name = None <EOL> if configuration . get ( '<STR_LIT>' , '<STR_LIT>' ) == '<STR_LIT>' : <EOL> auth_mechanism = db . extra_dejson . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> kerberos_service_name = db . extra_dejson . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> if auth_mechanism == '<STR_LIT>' : <EOL> logging . warning ( "<STR_LIT>" , <EOL> self . hiveserver2_conn_id ) <EOL> auth_mechanism = '<STR_LIT>' <EOL> from impala . dbapi import connect <EOL> return connect ( <EOL> host = db . host , <EOL> port = db . port , <EOL> auth_mechanism = auth_mechanism , <EOL> kerberos_service_name = kerberos_service_name , <EOL> user = db . login , <EOL> database = db . schema or '<STR_LIT:default>' ) <EOL> def get_results ( self , hql , schema = '<STR_LIT:default>' , arraysize = <NUM_LIT:1000> ) : <EOL> with self . get_conn ( ) as conn : <EOL> if isinstance ( hql , basestring ) : <EOL> hql = [ hql ] <EOL> results = { <EOL> '<STR_LIT:data>' : [ ] , <EOL> '<STR_LIT>' : [ ] , <EOL> } <EOL> for statement in hql : <EOL> with conn . cursor ( ) as cur : <EOL> cur . execute ( statement ) <EOL> records = cur . fetchall ( ) <EOL> if records : <EOL> results = { <EOL> '<STR_LIT:data>' : records , <EOL> '<STR_LIT>' : cur . description , <EOL> } <EOL> return results <EOL> def to_csv ( <EOL> self , <EOL> hql , <EOL> csv_filepath , <EOL> schema = '<STR_LIT:default>' , <EOL> delimiter = '<STR_LIT:U+002C>' , <EOL> lineterminator = '<STR_LIT:\r\n>' , <EOL> output_header = True ) : <EOL> schema = schema or '<STR_LIT:default>' <EOL> with self . get_conn ( ) as conn : <EOL> with conn . cursor ( ) as cur : <EOL> logging . info ( "<STR_LIT>" + hql ) <EOL> cur . execute ( hql ) <EOL> schema = cur . description <EOL> with open ( csv_filepath , '<STR_LIT:wb>' ) as f : <EOL> writer = csv . writer ( f , delimiter = delimiter , <EOL> lineterminator = lineterminator , encoding = '<STR_LIT:utf-8>' ) <EOL> if output_header : <EOL> writer . writerow ( [ c [ <NUM_LIT:0> ] <EOL> for c in cur . description ] ) <EOL> i = <NUM_LIT:0> <EOL> while True : <EOL> rows = [ row for row in cur . fetchmany ( ) if row ] <EOL> if not rows : <EOL> break <EOL> writer . writerows ( rows ) <EOL> i += len ( rows ) <EOL> logging . info ( "<STR_LIT>" . format ( i ) ) <EOL> logging . info ( "<STR_LIT>" . format ( i ) ) <EOL> def get_records ( self , hql , schema = '<STR_LIT:default>' ) : <EOL> """<STR_LIT>""" <EOL> return self . get_results ( hql , schema = schema ) [ '<STR_LIT:data>' ] <EOL> def get_pandas_df ( self , hql , schema = '<STR_LIT:default>' ) : <EOL> """<STR_LIT>""" <EOL> import pandas as pd <EOL> res = self . get_results ( hql , schema = schema ) <EOL> df = pd . DataFrame ( res [ '<STR_LIT:data>' ] ) <EOL> df . columns = [ c [ <NUM_LIT:0> ] for c in res [ '<STR_LIT>' ] ] <EOL> return df </s>
