<commit_before>from __future__ import print_function
from builtins import zip
from past.builtins import basestring
import csv
import logging
import re
import subprocess
import socket
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.utils import AirflowException
from airflow.hooks.base_hook import BaseHook
from airflow.utils import TemporaryDirectory
from airflow.configuration import conf


class HiveCliHook(BaseHook):
    """
    Simple wrapper around the hive CLI.

    It also supports the ``beeline``
    a lighter CLI that runs JDBC and is replacing the heavier
    traditional CLI. To enable ``beeline``, set the use_beeline param in the
    extra field of your connection as in ``{ "use_beeline": true }``

    Note that you can also set default hive CLI parameters using the
    ``hive_cli_params`` to be used in your connection as in
    ``{"hive_cli_params": "-hiveconf mapred.job.tracker=some.jobtracker:444"}``
    """

    def __init__(
            self,
            hive_cli_conn_id="hive_cli_default"):
        conn = self.get_connection(hive_cli_conn_id)
        self.hive_cli_params = conn.extra_dejson.get('hive_cli_params', '')
        self.use_beeline = conn.extra_dejson.get('use_beeline', False)
        self.conn = conn

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

        >>> hh = HiveCliHook()
        >>> result = hh.run_cli("USE airflow;")
        >>> ("OK" in result)
        True
        """
        conn = self.conn
        schema = schema or conn.schema
        if schema:
            hql = "USE {schema};\n{hql}".format(**locals())

        with TemporaryDirectory(prefix='airflow_hiveop_') as tmp_dir:
            with NamedTemporaryFile(dir=tmp_dir) as f:
                f.write(hql)
                f.flush()
                fname = f.name
                hive_bin = 'hive'
                cmd_extra = []
                if self.use_beeline:
                    hive_bin = 'beeline'
                    if conf.get('security','enabled'):
                        template = conn.extra_dejson.get('principal',"hive/_<EMAIL>")
                        template = template.replace("_HOST", socket.getfqdn())
                        jdbc_url = (
                            "jdbc:hive2://"
                            "{0}:{1}/{2}"
                            ";principal={3}"
                        ).format(conn.host, conn.port, conn.schema, template)
                    else:
                        jdbc_url = (
                            "jdbc:hive2://"
                            "{0}:{1}/{2}"
                            ";auth=noSasl"
                        ).format(conn.host, conn.port, conn.schema)

                    cmd_extra += ['-u', jdbc_url]
                    if conn.login:
                        cmd_extra += ['-n', conn.login]
                    if conn.password:
                        cmd_extra += ['-p', conn.password]
                    cmd_extra += ['-p', conn.login]
                hive_cmd = [hive_bin, '-f', fname] + cmd_extra
                if self.hive_cli_params:
                    hive_params_list = self.hive_cli_params.split()
                    hive_cmd.extend(hive_params_list)
                if verbose:
                    logging.info(" ".join(hive_cmd))
                sp = subprocess.Popen(
                    hive_cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    cwd=tmp_dir)
                self.sp = sp
                stdout = ''
                for line in iter(sp.stdout.readline, ''):
                    stdout += line
                    if verbose:
                        logging.info(line.strip())
                sp.wait()

                if sp.returncode:
                    raise AirflowException(stdout)

                return stdout


    def test_hql(self, hql):
        """
        Test an hql statement using the hive cli and EXPLAIN
<commit_msg>remove cmd_extra awkward addition<commit_after>from __future__ import print_function
from builtins import zip
from past.builtins import basestring
import csv
import logging
import re
import subprocess
import socket
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.utils import AirflowException
from airflow.hooks.base_hook import BaseHook
from airflow.utils import TemporaryDirectory
from airflow.configuration import conf


class HiveCliHook(BaseHook):
    """
    Simple wrapper around the hive CLI.

    It also supports the ``beeline``
    a lighter CLI that runs JDBC and is replacing the heavier
    traditional CLI. To enable ``beeline``, set the use_beeline param in the
    extra field of your connection as in ``{ "use_beeline": true }``

    Note that you can also set default hive CLI parameters using the
    ``hive_cli_params`` to be used in your connection as in
    ``{"hive_cli_params": "-hiveconf mapred.job.tracker=some.jobtracker:444"}``
    """

    def __init__(
            self,
            hive_cli_conn_id="hive_cli_default"):
        conn = self.get_connection(hive_cli_conn_id)
        self.hive_cli_params = conn.extra_dejson.get('hive_cli_params', '')
        self.use_beeline = conn.extra_dejson.get('use_beeline', False)
        self.conn = conn

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

        >>> hh = HiveCliHook()
        >>> result = hh.run_cli("USE airflow;")
        >>> ("OK" in result)
        True
        """
        conn = self.conn
        schema = schema or conn.schema
        if schema:
            hql = "USE {schema};\n{hql}".format(**locals())

        with TemporaryDirectory(prefix='airflow_hiveop_') as tmp_dir:
            with NamedTemporaryFile(dir=tmp_dir) as f:
                f.write(hql)
                f.flush()
                fname = f.name
                hive_bin = 'hive'
                cmd_extra = []

                if self.use_beeline:
                    hive_bin = 'beeline'
                    if conf.get('security','enabled'):
                        template = conn.extra_dejson.get('principal',"hive/_<EMAIL>")
                        template = template.replace("_HOST", socket.getfqdn())
                        jdbc_url = (
                            "jdbc:hive2://"
                            "{0}:{1}/{2}"
                            ";principal={3}"
                        ).format(conn.host, conn.port, conn.schema, template)
                    else:
                        jdbc_url = (
                            "jdbc:hive2://"
                            "{0}:{1}/{2}"
                            ";auth=noSasl"
                        ).format(conn.host, conn.port, conn.schema)

                    cmd_extra += ['-u', jdbc_url]
                    if conn.login:
                        cmd_extra += ['-n', conn.login]
                    if conn.password:
                        cmd_extra += ['-p', conn.password]

                hive_cmd = [hive_bin, '-f', fname] + cmd_extra

                if self.hive_cli_params:
                    hive_params_list = self.hive_cli_params.split()
                    hive_cmd.extend(hive_params_list)
                if verbose:
                    logging.info(" ".join(hive_cmd))
                sp = subprocess.Popen(
                    hive_cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    cwd=tmp_dir)
                self.sp = sp
                stdout = ''
                for line in iter(sp.stdout.readline, ''):
                    stdout += line
                    if verbose:
                        logging.info(line.strip())
                sp.wait()

                if sp.returncode:
                    raise AirflowException(stdout)

                return stdout


    def test_hql(self, hql):
        """
        Test an hql statement using the hive cli and EXPLAIN

{'fraction_non_alphanumeric': 0.09376875975507264, 'fraction_numerical': 0.003121623244086925, 'mean_word_length': 1.5575683143997543, 'pattern_counts': {'":': 4, '<': 5, '<?xml version=': 0, '>': 23, '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': '3072162', 'n_tokens_mistral': 2400, 'n_tokens_neox': 2245, 'n_words': 568}