<commit_before>import csv
import logging
import json
import subprocess
from tempfile import NamedTemporaryFile

from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol
from hive_service import ThriftHive
import pyhs2

from airflow.hooks.base_hook import BaseHook


class HiveCliHook(BaseHook):
    '''
    Simple wrapper around the hive CLI
    '''
    def __init__(
            self,
            hive_cli_conn_id="hive_cli_default"
            ):
        conn = self.get_connection(hive_cli_conn_id)
        self.hive_cli_params = ""
        try:
            self.hive_cli_params = json.loads(conn.extra)['hive_cli_params']
        except:
            pass

    def run_cli(self, hql, schema=None):
        '''
        Run an hql statement using the hive cli

        >>> hh = HiveCliHook()
        >>> hh.run_cli("USE airflow;")
        '''
        if schema:
            hql = "USE {schema};\n{hql}".format(**locals())

        f = NamedTemporaryFile()
        f.write(hql)
        f.flush()
        sp = subprocess.Popen(
            "hive -f {f.name} {self.hive_cli_params}".format(**locals()),
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT)
        all_err = ''
        self.sp = sp
        for line in iter(sp.stdout.readline, ''):
            logging.info(line.strip())
        sp.wait()
        f.close()

        if sp.returncode:
            raise Exception(all_err)

    def kill(self):
        if hasattr(self, 'sp'):
            if self.sp.poll() is None:
                print("Killing the Hive job")
                self.sp.kill()


class HiveMetastoreHook(BaseHook):
    '''
    Wrapper to interact with the Hive Metastore
    '''
    def __init__(self, metastore_conn_id='metastore_default'):
        self.metastore_conn = self.get_connection(metastore_conn_id)<commit_msg>[hive_hook] add with statement to ensure nice cleaning in case of an error<commit_after>import csv
import logging
import json
import subprocess
from tempfile import NamedTemporaryFile


from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol
from hive_service import ThriftHive
import pyhs2

from airflow.hooks.base_hook import BaseHook


class HiveCliHook(BaseHook):
    '''
    Simple wrapper around the hive CLI
    '''
    def __init__(
            self,
            hive_cli_conn_id="hive_cli_default"
            ):
        conn = self.get_connection(hive_cli_conn_id)
        self.hive_cli_params = ""
        try:
            self.hive_cli_params = json.loads(conn.extra)['hive_cli_params']
        except:
            pass

    def run_cli(self, hql, schema=None):
        '''
        Run an hql statement using the hive cli

        >>> hh = HiveCliHook()
        >>> hh.run_cli("USE airflow;")
        '''
        if schema:
            hql = "USE {schema};\n{hql}".format(**locals())

        with NamedTemporaryFile() as f:
            f.write(hql)
            f.flush()
            fname = f.name
            hive_cmd = ['hive' , '-f' , fname]
            if self.hive_cli_params:
                hive_params_list = self.hive_cli_params.split()
                hive_cmd.extend(hive_params_list)
            sp = subprocess.Popen(
                hive_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT)
            all_err = ''
            self.sp = sp
            for line in iter(sp.stdout.readline, ''):
                logging.info(line.strip())
            sp.wait()

            if sp.returncode:
                raise Exception(all_err)

    def kill(self):
        if hasattr(self, 'sp'):
            if self.sp.poll() is None:
                print("Killing the Hive job")
                self.sp.kill()


class HiveMetastoreHook(BaseHook):
    '''
    Wrapper to interact with the Hive Metastore
    '''
    def __init__(self, metastore_conn_id='metastore_default'):
        self.metastore_conn = self.get_connection(metastore_conn_id)
{'fraction_non_alphanumeric': 0.09763895720609936, 'fraction_numerical': 0.0004918839153959665, 'mean_word_length': 2.055597295266717, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7120144', 'n_tokens_mistral': 1281, 'n_tokens_neox': 1233, 'n_words': 279}