# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import datetime
import os
import unittest

from unittest import mock
import nose

from airflow import DAG, configuration, operators
from airflow.models import TaskInstance
from airflow.operators.hive_operator import HiveOperator
from airflow.utils import timezone


DEFAULT_DATE = datetime.datetime(2015, 1, 1)
DEFAULT_DATE_ISO = DEFAULT_DATE.isoformat()
DEFAULT_DATE_DS = DEFAULT_DATE_ISO[:10]


class TestHiveEnvironment(unittest.TestCase):

    def setUp(self):
        args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
        dag = DAG('test_dag_id', default_args=args)
        self.dag = dag
        self.hql = """
        USE airflow;
        DROP TABLE IF EXISTS static_babynames_partitioned;
        CREATE TABLE IF NOT EXISTS static_babynames_partitioned (
            state string,
            year string,
            name string,
            gender string,
            num int)
        PARTITIONED BY (ds string);
        INSERT OVERWRITE TABLE static_babynames_partitioned
            PARTITION(ds='{{ ds }}')
        SELECT state, year, name, gender, num FROM static_babynames;
        """


class TestHiveCli(unittest.TestCase):

    def setUp(self):
        self.nondefault_schema = "nondefault"
        os.environ["AIRFLOW__CORE__SECURITY"] = "kerberos"

    def tearDown(self):
        del os.environ["AIRFLOW__CORE__SECURITY"]

    def test_get_proxy_user_value(self):
        from airflow.hooks.hive_hooks import HiveCliHook

        hook = HiveCliHook()
        returner = mock.MagicMock()
        returner.extra_dejson = {'proxy_user': 'a_user_proxy'}
        hook.use_beeline = True
        hook.conn = returner

        # Run
        result = hook._prepare_cli_cmd()

        # Verify
        self.assertIn('hive.server2.proxy.user=a_user_proxy', result[2])


class HiveOperatorConfigTest(TestHiveEnvironment):

    def test_hive_airflow_default_config_queue(self):
        t = HiveOperator(
            task_id='test_default_config_queue',
            hql=self.hql,
            mapred_queue_priority='HIGH',
            mapred_job_name='airflow.test_default_config_queue',
            dag=self.dag)

        # just check that the correct default value in test_default.cfg is used
        test_config_hive_mapred_queue = configuration.conf.get(
            'hive',
            'default_hive_mapred_queue'
        )
        self.assertEqual(t.get_hook().mapred_queue, test_config_hive_mapred_queue)

    def test_hive_airflow_default_config_queue_override(self):
        specific_mapred_queue = 'default'
        t = HiveOperator(
            task_id='test_default_config_queue',
            hql=self.hql,
            mapred_queue=specific_mapred_queue,
            mapred_queue_priority='HIGH',
            mapred_job_name='airflow.test_default_config_queue',
            dag=self.dag)

        self.assertEqual(t.get_hook().mapred_queue, specific_mapred_queue)


class HiveOperatorTest(TestHiveEnvironment):

    def test_hiveconf_jinja_translate(self):
        hql = "SELECT ${num_col} FROM ${hiveconf:table};"
        t = HiveOperator(
            hiveconf_jinja_translate=True,
            task_id='dry_run_basic_hql', hql=hql, dag=self.dag)
        t.prepare_template()
        self.assertEqual(t.hql, "SELECT {{ num_col }} FROM {{ table }};")

    def test_hiveconf(self):
        hql = "SELECT * FROM ${hiveconf:table} PARTITION (${hiveconf:day});"
        t = HiveOperator(
            hiveconfs={'table': 'static_babynames', 'day': '{{ ds }}'},
            task_id='dry_run_basic_hql', hql=hql, dag=self.dag)
        t.prepare_template()
        self.assertEqual(
            t.hql,
            "SELECT * FROM ${hiveconf:table} PARTITION (${hiveconf:day});")

    @mock.patch('airflow.operators.hive_operator.HiveOperator.get_hook')
    def test_mapred_job_name(self, mock_get_hook):
        mock_hook = mock.MagicMock()
        mock_get_hook.return_value = mock_hook
        t = HiveOperator(
            task_id='test_mapred_job_name',
            hql=self.hql,
            dag=self.dag)

        fake_execution_date = timezone.datetime(2018, 6, 19)
        fake_ti = TaskInstance(task=t, execution_date=fake_execution_date)
        fake_ti.hostname = 'fake_hostname'
        fake_context = {'ti': fake_ti}

        t.execute(fake_context)
        self.assertEqual(
            "Airflow HiveOperator task for {}.{}.{}.{}"
            .format(fake_ti.hostname,
                    self.dag.dag_id, t.task_id,
                    fake_execution_date.isoformat()), mock_hook.mapred_job_name)


if 'AIRFLOW_RUNALL_TESTS' in os.environ:

    import airflow.hooks.hive_hooks
    import airflow.operators.presto_to_mysql

    class TestHivePresto(TestHiveEnvironment):

        def test_hive(self):
            t = HiveOperator(
                task_id='basic_hql', hql=self.hql, dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_hive_queues(self):
            t = HiveOperator(
                task_id='test_hive_queues', hql=self.hql,
                mapred_queue='default', mapred_queue_priority='HIGH',
                mapred_job_name='airflow.test_hive_queues',
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_hive_dryrun(self):
            t = HiveOperator(
                task_id='dry_run_basic_hql', hql=self.hql, dag=self.dag)
            t.dry_run()

        def test_beeline(self):
            t = HiveOperator(
                task_id='beeline_hql', hive_cli_conn_id='hive_cli_default',
                hql=self.hql, dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_presto(self):
            sql = """
            SELECT count(1) FROM airflow.static_babynames_partitioned;
            """
            t = operators.presto_check_operator.PrestoCheckOperator(
                task_id='presto_check', sql=sql, dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_presto_to_mysql(self):
            t = operators.presto_to_mysql.PrestoToMySqlTransfer(
                task_id='presto_to_mysql_check',
                sql="""
                SELECT name, count(*) as ccount
                FROM airflow.static_babynames
                GROUP BY name
                """,
                mysql_table='test_static_babynames',
                mysql_preoperator='TRUNCATE TABLE test_static_babynames;',
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_hdfs_sensor(self):
            t = operators.sensors.HdfsSensor(
                task_id='hdfs_sensor_check',
                filepath='hdfs://user/hive/warehouse/airflow.db/static_babynames',
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_webhdfs_sensor(self):
            t = operators.sensors.WebHdfsSensor(
                task_id='webhdfs_sensor_check',
                filepath='hdfs://user/hive/warehouse/airflow.db/static_babynames',
                timeout=120,
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_sql_sensor(self):
            t = operators.sensors.SqlSensor(
                task_id='hdfs_sensor_check',
                conn_id='presto_default',
                sql="SELECT 'x' FROM airflow.static_babynames LIMIT 1;",
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_hive_stats(self):
            t = operators.hive_stats_operator.HiveStatsCollectionOperator(
                task_id='hive_stats_check',
                table="airflow.static_babynames_partitioned",
                partition={'ds': DEFAULT_DATE_DS},
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_named_hive_partition_sensor(self):
            t = operators.sensors.NamedHivePartitionSensor(
                task_id='hive_partition_check',
                partition_names=[
                    "airflow.static_babynames_partitioned/ds={{ds}}"
                ],
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_named_hive_partition_sensor_succeeds_on_multiple_partitions(self):
            t = operators.sensors.NamedHivePartitionSensor(
                task_id='hive_partition_check',
                partition_names=[
                    "airflow.static_babynames_partitioned/ds={{ds}}",
                    "airflow.static_babynames_partitioned/ds={{ds}}"
                ],
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_named_hive_partition_sensor_parses_partitions_with_periods(self):
            t = operators.sensors.NamedHivePartitionSensor.parse_partition_name(
                partition="schema.table/part1=this.can.be.an.issue/part2=ok")
            self.assertEqual(t[0], "schema")
            self.assertEqual(t[1], "table")
            self.assertEqual(t[2], "part1=this.can.be.an.issue/part2=this_should_be_ok")

        @nose.tools.raises(airflow.exceptions.AirflowSensorTimeout)
        def test_named_hive_partition_sensor_times_out_on_nonexistent_partition(self):
            t = operators.sensors.NamedHivePartitionSensor(
                task_id='hive_partition_check',
                partition_names=[
                    "airflow.static_babynames_partitioned/ds={{ds}}",
                    "airflow.static_babynames_partitioned/ds=nonexistent"
                ],
                poke_interval=0.1,
                timeout=1,
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_hive_partition_sensor(self):
            t = operators.sensors.HivePartitionSensor(
                task_id='hive_partition_check',
                table='airflow.static_babynames_partitioned',
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_hive_metastore_sql_sensor(self):
            t = operators.sensors.MetastorePartitionSensor(
                task_id='hive_partition_check',
                table='airflow.static_babynames_partitioned',
                partition_name='ds={}'.format(DEFAULT_DATE_DS),
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_hive2samba(self):
            t = operators.hive_to_samba_operator.Hive2SambaOperator(
                task_id='hive2samba_check',
                samba_conn_id='tableau_samba',
                hql="SELECT * FROM airflow.static_babynames LIMIT 10000",
                destination_filepath='test_airflow.csv',
                dag=self.dag)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

        def test_hive_to_mysql(self):
            t = operators.hive_to_mysql.HiveToMySqlTransfer(
                mysql_conn_id='airflow_db',
                task_id='hive_to_mysql_check',
                create=True,
                sql="""
                SELECT name
                FROM airflow.static_babynames
                LIMIT 100
                """,
                mysql_table='test_static_babynames',
                mysql_preoperator=[
                    'DROP TABLE IF EXISTS test_static_babynames;',
                    'CREATE TABLE test_static_babynames (name VARCHAR(500))',
                ],
                dag=self.dag)
            t.clear(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
                  ignore_ti_state=True)

[(33, 'datetime.datetime', 'datetime.datetime', (['(2015)', '(1)', '(1)'], {}), False, 'import datetime\n'), (135, 'unittest.mock.patch', 'mock.patch', (['"""airflow.operators.hive_operator.HiveOperator.get_hook"""'], {}), False, 'from unittest import mock\n'), (42, 'airflow.DAG', 'DAG', (['"""test_dag_id"""'], {'default_args': 'args'}), False, 'from airflow import DAG, configuration, operators\n'), (72, 'airflow.hooks.hive_hooks.HiveCliHook', 'HiveCliHook', ([], {}), False, 'from airflow.hooks.hive_hooks import HiveCliHook\n'), (73, 'unittest.mock.MagicMock', 'mock.MagicMock', ([], {}), False, 'from unittest import mock\n'), (88, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'task_id': '"""test_default_config_queue"""', 'hql': 'self.hql', 'mapred_queue_priority': '"""HIGH"""', 'mapred_job_name': '"""airflow.test_default_config_queue"""', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (96, 'airflow.configuration.conf.get', 'configuration.conf.get', (['"""hive"""', '"""default_hive_mapred_queue"""'], {}), False, 'from airflow import DAG, configuration, operators\n'), (104, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'task_id': '"""test_default_config_queue"""', 'hql': 'self.hql', 'mapred_queue': 'specific_mapred_queue', 'mapred_queue_priority': '"""HIGH"""', 'mapred_job_name': '"""airflow.test_default_config_queue"""', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (119, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'hiveconf_jinja_translate': '(True)', 'task_id': '"""dry_run_basic_hql"""', 'hql': 'hql', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (127, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'hiveconfs': "{'table': 'static_babynames', 'day': '{{ ds }}'}", 'task_id': '"""dry_run_basic_hql"""', 'hql': 'hql', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (137, 'unittest.mock.MagicMock', 'mock.MagicMock', ([], {}), False, 'from unittest import mock\n'), (139, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'task_id': '"""test_mapred_job_name"""', 'hql': 'self.hql', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (144, 'airflow.utils.timezone.datetime', 'timezone.datetime', (['(2018)', '(6)', '(19)'], {}), False, 'from airflow.utils import timezone\n'), (145, 'airflow.models.TaskInstance', 'TaskInstance', ([], {'task': 't', 'execution_date': 'fake_execution_date'}), False, 'from airflow.models import TaskInstance\n'), (277, 'nose.tools.raises', 'nose.tools.raises', (['airflow.exceptions.AirflowSensorTimeout'], {}), False, 'import nose\n'), (165, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'task_id': '"""basic_hql"""', 'hql': 'self.hql', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (171, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'task_id': '"""test_hive_queues"""', 'hql': 'self.hql', 'mapred_queue': '"""default"""', 'mapred_queue_priority': '"""HIGH"""', 'mapred_job_name': '"""airflow.test_hive_queues"""', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (180, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'task_id': '"""dry_run_basic_hql"""', 'hql': 'self.hql', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (185, 'airflow.operators.hive_operator.HiveOperator', 'HiveOperator', ([], {'task_id': '"""beeline_hql"""', 'hive_cli_conn_id': '"""hive_cli_default"""', 'hql': 'self.hql', 'dag': 'self.dag'}), False, 'from airflow.operators.hive_operator import HiveOperator\n'), (195, 'airflow.operators.presto_check_operator.PrestoCheckOperator', 'operators.presto_check_operator.PrestoCheckOperator', ([], {'task_id': '"""presto_check"""', 'sql': 'sql', 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (201, 'airflow.operators.presto_to_mysql.PrestoToMySqlTransfer', 'operators.presto_to_mysql.PrestoToMySqlTransfer', ([], {'task_id': '"""presto_to_mysql_check"""', 'sql': '"""\n                SELECT name, count(*) as ccount\n                FROM airflow.static_babynames\n                GROUP BY name\n                """', 'mysql_table': '"""test_static_babynames"""', 'mysql_preoperator': '"""TRUNCATE TABLE test_static_babynames;"""', 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (215, 'airflow.operators.sensors.HdfsSensor', 'operators.sensors.HdfsSensor', ([], {'task_id': '"""hdfs_sensor_check"""', 'filepath': '"""hdfs://user/hive/warehouse/airflow.db/static_babynames"""', 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (223, 'airflow.operators.sensors.WebHdfsSensor', 'operators.sensors.WebHdfsSensor', ([], {'task_id': '"""webhdfs_sensor_check"""', 'filepath': '"""hdfs://user/hive/warehouse/airflow.db/static_babynames"""', 'timeout': '(120)', 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (232, 'airflow.operators.sensors.SqlSensor', 'operators.sensors.SqlSensor', ([], {'task_id': '"""hdfs_sensor_check"""', 'conn_id': '"""presto_default"""', 'sql': '"""SELECT \'x\' FROM airflow.static_babynames LIMIT 1;"""', 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (241, 'airflow.operators.hive_stats_operator.HiveStatsCollectionOperator', 'operators.hive_stats_operator.HiveStatsCollectionOperator', ([], {'task_id': '"""hive_stats_check"""', 'table': '"""airflow.static_babynames_partitioned"""', 'partition': "{'ds': DEFAULT_DATE_DS}", 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (250, 'airflow.operators.sensors.NamedHivePartitionSensor', 'operators.sensors.NamedHivePartitionSensor', ([], {'task_id': '"""hive_partition_check"""', 'partition_names': "['airflow.static_babynames_partitioned/ds={{ds}}']", 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (260, 'airflow.operators.sensors.NamedHivePartitionSensor', 'operators.sensors.NamedHivePartitionSensor', ([], {'task_id': '"""hive_partition_check"""', 'partition_names': "['airflow.static_babynames_partitioned/ds={{ds}}',\n    'airflow.static_babynames_partitioned/ds={{ds}}']", 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (271, 'airflow.operators.sensors.NamedHivePartitionSensor.parse_partition_name', 'operators.sensors.NamedHivePartitionSensor.parse_partition_name', ([], {'partition': '"""schema.table/part1=this.can.be.an.issue/part2=ok"""'}), False, 'from airflow import DAG, configuration, operators\n'), (279, 'airflow.operators.sensors.NamedHivePartitionSensor', 'operators.sensors.NamedHivePartitionSensor', ([], {'task_id': '"""hive_partition_check"""', 'partition_names': "['airflow.static_babynames_partitioned/ds={{ds}}',\n    'airflow.static_babynames_partitioned/ds=nonexistent']", 'poke_interval': '(0.1)', 'timeout': '(1)', 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (292, 'airflow.operators.sensors.HivePartitionSensor', 'operators.sensors.HivePartitionSensor', ([], {'task_id': '"""hive_partition_check"""', 'table': '"""airflow.static_babynames_partitioned"""', 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (309, 'airflow.operators.hive_to_samba_operator.Hive2SambaOperator', 'operators.hive_to_samba_operator.Hive2SambaOperator', ([], {'task_id': '"""hive2samba_check"""', 'samba_conn_id': '"""tableau_samba"""', 'hql': '"""SELECT * FROM airflow.static_babynames LIMIT 10000"""', 'destination_filepath': '"""test_airflow.csv"""', 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n'), (319, 'airflow.operators.hive_to_mysql.HiveToMySqlTransfer', 'operators.hive_to_mysql.HiveToMySqlTransfer', ([], {'mysql_conn_id': '"""airflow_db"""', 'task_id': '"""hive_to_mysql_check"""', 'create': '(True)', 'sql': '"""\n                SELECT name\n                FROM airflow.static_babynames\n                LIMIT 100\n                """', 'mysql_table': '"""test_static_babynames"""', 'mysql_preoperator': "['DROP TABLE IF EXISTS test_static_babynames;',\n    'CREATE TABLE test_static_babynames (name VARCHAR(500))']", 'dag': 'self.dag'}), False, 'from airflow import DAG, configuration, operators\n')]