[AIRFLOW-2463] Make task instance context available for hive queries

UPDATING.md

@@ -40,6 +40,7 @@ There are five roles created for Airflow by default: Admin, User, Op, Viewer, an
 - Airflow dag home page is now `/home` (instead of `/admin`).
 - All ModelViews in Flask-AppBuilder follow a different pattern from Flask-Admin. The `/admin` part of the url path will no longer exist. For example: `/admin/connection` becomes `/connection/list`, `/admin/connection/new` becomes `/connection/add`, `/admin/connection/edit` becomes `/connection/edit`, etc.
 - Due to security concerns, the new webserver will no longer support the features in the `Data Profiling` menu of old UI, including `Ad Hoc Query`, `Charts`, and `Known Events`.
+- HiveServer2Hook.get_results() always returns a list of tuples, even when a single column is queried, as per Python API 2.
 
 ### airflow.contrib.sensors.hdfs_sensors renamed to airflow.contrib.sensors.hdfs_sensor
 
@@ -72,7 +73,7 @@ supported and will be removed entirely in Airflow 2.0
 With Airflow 1.9 or lower, Unload operation always included header row. In order to include header row,
 we need to turn off parallel unload. It is preferred to perform unload operation using all nodes so that it is
 faster for larger tables. So, parameter called `include_header` is added and default is set to False.
-Header row will be added only if this parameter is set True and also in that case parallel will be automatically turned off (`PARALLEL OFF`)  
+Header row will be added only if this parameter is set True and also in that case parallel will be automatically turned off (`PARALLEL OFF`)
 
 ### Google cloud connection string
 

airflow/hooks/hive_hooks.py

@@ -21,30 +21,40 @@
 
 import contextlib
 import os
-
-from six.moves import zip
-from past.builtins import basestring, unicode
-
-import unicodecsv as csv
 import re
-import six
 import subprocess
 import time
 from collections import OrderedDict
 from tempfile import NamedTemporaryFile
+
 import hmsclient
+import six
+import unicodecsv as csv
+from past.builtins import basestring
+from past.builtins import unicode
+from six.moves import zip
 
-from airflow import configuration as conf
+import airflow.security.utils as utils
+from airflow import configuration
 from airflow.exceptions import AirflowException
 from airflow.hooks.base_hook import BaseHook
-from airflow.utils.helpers import as_flattened_list
 from airflow.utils.file import TemporaryDirectory
-from airflow import configuration
-import airflow.security.utils as utils
+from airflow.utils.helpers import as_flattened_list
+from airflow.utils.operator_helpers import AIRFLOW_VAR_NAME_FORMAT_MAPPING
 
 HIVE_QUEUE_PRIORITIES = ['VERY_HIGH', 'HIGH', 'NORMAL', 'LOW', 'VERY_LOW']
 
 
+def get_context_from_env_var():
+    """
+    Extract context from env variable, e.g. dag_id, task_id and execution_date,
+    so that they can be used inside BashOperator and PythonOperator.
+    :return: The context of interest.
+    """
+    return {format_map['default']: os.environ.get(format_map['env_var_format'], '')
+            for format_map in AIRFLOW_VAR_NAME_FORMAT_MAPPING.values()}
+
+
 class HiveCliHook(BaseHook):
     """Simple wrapper around the hive CLI.
 
@@ -92,8 +102,8 @@ def __init__(
                     "Invalid Mapred Queue Priority.  Valid values are: "
                     "{}".format(', '.join(HIVE_QUEUE_PRIORITIES)))
 
-        self.mapred_queue = mapred_queue or conf.get('hive',
-                                                     'default_hive_mapred_queue')
+        self.mapred_queue = mapred_queue or configuration.get('hive',
+                                                              'default_hive_mapred_queue')
         self.mapred_queue_priority = mapred_queue_priority
         self.mapred_job_name = mapred_job_name
 
@@ -126,6 +136,7 @@ def _prepare_cli_cmd(self):
                 jdbc_url += ";auth=" + self.auth
 
             jdbc_url = jdbc_url.format(**locals())
+            jdbc_url = '"{}"'.format(jdbc_url)
 
             cmd_extra += ['-u', jdbc_url]
             if conn.login:
@@ -184,10 +195,15 @@ def run_cli(self, hql, schema=None, verbose=True, hive_conf=None):
 
         with TemporaryDirectory(prefix='airflow_hiveop_') as tmp_dir:
             with NamedTemporaryFile(dir=tmp_dir) as f:
+                hql = hql + '\n'
                 f.write(hql.encode('UTF-8'))
                 f.flush()
                 hive_cmd = self._prepare_cli_cmd()
-                hive_conf_params = self._prepare_hiveconf(hive_conf)
+                env_context = get_context_from_env_var()
+                # Only extend the hive_conf if it is defined.
+                if hive_conf:
+                    env_context.update(hive_conf)
+                hive_conf_params = self._prepare_hiveconf(env_context)
                 if self.mapred_queue:
                     hive_conf_params.extend(
                         ['-hiveconf',
@@ -772,20 +788,29 @@ def get_conn(self, schema=None):
             username=db.login or username,
             database=schema or db.schema or 'default')
 
-    def _get_results(self, hql, schema='default', fetch_size=None):
+    def _get_results(self, hql, schema='default', fetch_size=None, hive_conf=None):
         from pyhive.exc import ProgrammingError
         if isinstance(hql, basestring):
             hql = [hql]
         previous_description = None
         with contextlib.closing(self.get_conn(schema)) as conn, \
                 contextlib.closing(conn.cursor()) as cur:
             cur.arraysize = fetch_size or 1000
+
+            env_context = get_context_from_env_var()
+            if hive_conf:
+                env_context.update(hive_conf)
+            for k, v in env_context.items():
+                cur.execute("set {}={}".format(k, v))
+
             for statement in hql:
                 cur.execute(statement)
                 # we only get results of statements that returns
                 lowered_statement = statement.lower().strip()
                 if (lowered_statement.startswith('select') or
-                   lowered_statement.startswith('with')):
+                    lowered_statement.startswith('with') or
+                    (lowered_statement.startswith('set') and
+                     '=' not in lowered_statement)):
                     description = [c for c in cur.description]
                     if previous_description and previous_description != description:
                         message = '''The statements are producing different descriptions:
@@ -805,8 +830,17 @@ def _get_results(self, hql, schema='default', fetch_size=None):
                     except ProgrammingError:
                         self.log.debug("get_results returned no records")
 
-    def get_results(self, hql, schema='default', fetch_size=None):
-        results_iter = self._get_results(hql, schema, fetch_size=fetch_size)
+    def get_results(self, hql, schema='default', fetch_size=None, hive_conf=None):
+        """
+        Get results of the provided hql in target schema.
+        :param hql: hql to be executed.
+        :param schema: target schema, default to 'default'.
+        :param fetch_size max size of result to fetch.
+        :param hive_conf: hive_conf to execute alone with the hql.
+        :return: results of hql execution.
+        """
+        results_iter = self._get_results(hql, schema,
+                                         fetch_size=fetch_size, hive_conf=hive_conf)
         header = next(results_iter)
         results = {
             'data': list(results_iter),
@@ -822,9 +856,23 @@ def to_csv(
             delimiter=',',
             lineterminator='\r\n',
             output_header=True,
-            fetch_size=1000):
-
-        results_iter = self._get_results(hql, schema, fetch_size=fetch_size)
+            fetch_size=1000,
+            hive_conf=None):
+        """
+        Execute hql in target schema and write results to a csv file.
+        :param hql: hql to be executed.
+        :param csv_filepath: filepath of csv to write results into.
+        :param schema: target schema, , default to 'default'.
+        :param delimiter: delimiter of the csv file.
+        :param lineterminator: lineterminator of the csv file.
+        :param output_header: header of the csv file.
+        :param fetch_size: number of result rows to write into the csv file.
+        :param hive_conf: hive_conf to execute alone with the hql.
+        :return:
+        """
+
+        results_iter = self._get_results(hql, schema,
+                                         fetch_size=fetch_size, hive_conf=hive_conf)
         header = next(results_iter)
         message = None
 

airflow/operators/bash_operator.py

@@ -18,16 +18,18 @@
 # under the License.
 
 
-from builtins import bytes
 import os
 import signal
 from subprocess import Popen, STDOUT, PIPE
 from tempfile import gettempdir, NamedTemporaryFile
 
+from builtins import bytes
+
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.utils.decorators import apply_defaults
 from airflow.utils.file import TemporaryDirectory
+from airflow.utils.operator_helpers import context_to_airflow_vars
 
 
 class BashOperator(BaseOperator):
@@ -73,6 +75,16 @@ def execute(self, context):
         """
         self.log.info("Tmp dir root location: \n %s", gettempdir())
 
+        # Prepare env for child process.
+        if self.env is None:
+            self.env = os.environ.copy()
+        airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True)
+        self.log.info("Exporting the following env vars:\n" +
+                      '\n'.join(["{}={}".format(k, v)
+                                 for k, v in
+                                 airflow_context_vars.items()]))
+        self.env.update(airflow_context_vars)
+
         self.lineage_data = self.bash_command
 
         with TemporaryDirectory(prefix='airflowtmp') as tmp_dir:

airflow/operators/hive_to_mysql.py

@@ -17,12 +17,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from tempfile import NamedTemporaryFile
+
 from airflow.hooks.hive_hooks import HiveServer2Hook
 from airflow.hooks.mysql_hook import MySqlHook
 from airflow.models import BaseOperator
 from airflow.utils.decorators import apply_defaults
-
-from tempfile import NamedTemporaryFile
+from airflow.utils.operator_helpers import context_to_airflow_vars
 
 
 class HiveToMySqlTransfer(BaseOperator):
@@ -88,7 +89,8 @@ def execute(self, context):
         if self.bulk_load:
             tmpfile = NamedTemporaryFile()
             hive.to_csv(self.sql, tmpfile.name, delimiter='\t',
-                        lineterminator='\n', output_header=False)
+                        lineterminator='\n', output_header=False,
+                        hive_conf=context_to_airflow_vars(context))
         else:
             results = hive.get_records(self.sql)
 

airflow/operators/hive_to_samba_operator.py

@@ -23,6 +23,7 @@
 from airflow.hooks.samba_hook import SambaHook
 from airflow.models import BaseOperator
 from airflow.utils.decorators import apply_defaults
+from airflow.utils.operator_helpers import context_to_airflow_vars
 
 
 class Hive2SambaOperator(BaseOperator):
@@ -60,6 +61,7 @@ def execute(self, context):
         hive = HiveServer2Hook(hiveserver2_conn_id=self.hiveserver2_conn_id)
         tmpfile = tempfile.NamedTemporaryFile()
         self.log.info("Fetching file from Hive")
-        hive.to_csv(hql=self.hql, csv_filepath=tmpfile.name)
+        hive.to_csv(hql=self.hql, csv_filepath=tmpfile.name,
+                    hive_conf=context_to_airflow_vars(context))
         self.log.info("Pushing to samba")
         samba.push_from_local(self.destination_filepath, tmpfile.name)

airflow/operators/python_operator.py

@@ -17,21 +17,22 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from builtins import str
-import dill
 import inspect
 import os
 import pickle
 import subprocess
 import sys
 import types
+from textwrap import dedent
+
+import dill
+from builtins import str
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator, SkipMixin
 from airflow.utils.decorators import apply_defaults
 from airflow.utils.file import TemporaryDirectory
-
-from textwrap import dedent
+from airflow.utils.operator_helpers import context_to_airflow_vars
 
 
 class PythonOperator(BaseOperator):
@@ -91,6 +92,13 @@ def __init__(
             self.template_ext = templates_exts
 
     def execute(self, context):
+        # Export context to make it available for callables to use.
+        airflow_context_vars = context_to_airflow_vars(context, in_env_var_format=True)
+        self.log.info("Exporting the following env vars:\n" +
+                      '\n'.join(["{}={}".format(k, v)
+                                 for k, v in airflow_context_vars.items()]))
+        os.environ.update(airflow_context_vars)
+
         if self.provide_context:
             context.update(self.op_kwargs)
             context['templates_dict'] = self.templates_dict

airflow/utils/operator_helpers.py

@@ -7,9 +7,9 @@
 # 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
@@ -18,32 +18,49 @@
 # under the License.
 #
 
+AIRFLOW_VAR_NAME_FORMAT_MAPPING = {
+    'AIRFLOW_CONTEXT_DAG_ID': {'default': 'airflow.ctx.dag_id',
+                               'env_var_format': 'AIRFLOW_CTX_DAG_ID'},
+    'AIRFLOW_CONTEXT_TASK_ID': {'default': 'airflow.ctx.task_id',
+                                'env_var_format': 'AIRFLOW_CTX_TASK_ID'},
+    'AIRFLOW_CONTEXT_EXECUTION_DATE': {'default': 'airflow.ctx.execution_date',
+                                       'env_var_format': 'AIRFLOW_CTX_EXECUTION_DATE'},
+    'AIRFLOW_CONTEXT_DAG_RUN_ID': {'default': 'airflow.ctx.dag_run_id',
+                                   'env_var_format': 'AIRFLOW_CTX_DAG_RUN_ID'}
+}
+
 
-def context_to_airflow_vars(context):
+def context_to_airflow_vars(context, in_env_var_format=False):
     """
     Given a context, this function provides a dictionary of values that can be used to
     externally reconstruct relations between dags, dag_runs, tasks and task_instances.
+    Default to abc.def.ghi format and can be made to ABC_DEF_GHI format if
+    in_env_var_format is set to True.
 
-    :param context: The context for the task_instance of interest
+    :param context: The context for the task_instance of interest.
     :type context: dict
+    :param in_env_var_format: If returned vars should be in ABC_DEF_GHI format.
+    :type in_env_var_format: bool
+    :return task_instance context as dict.
     """
-    params = {}
-    dag = context.get('dag')
-    if dag and dag.dag_id:
-        params['airflow.ctx.dag.dag_id'] = dag.dag_id
-
-    dag_run = context.get('dag_run')
-    if dag_run and dag_run.execution_date:
-        params['airflow.ctx.dag_run.execution_date'] = dag_run.execution_date.isoformat()
-
-    task = context.get('task')
-    if task and task.task_id:
-        params['airflow.ctx.task.task_id'] = task.task_id
-
+    params = dict()
+    if in_env_var_format:
+        name_format = 'env_var_format'
+    else:
+        name_format = 'default'
     task_instance = context.get('task_instance')
+    if task_instance and task_instance.dag_id:
+        params[AIRFLOW_VAR_NAME_FORMAT_MAPPING['AIRFLOW_CONTEXT_DAG_ID'][
+            name_format]] = task_instance.dag_id
+    if task_instance and task_instance.task_id:
+        params[AIRFLOW_VAR_NAME_FORMAT_MAPPING['AIRFLOW_CONTEXT_TASK_ID'][
+            name_format]] = task_instance.task_id
     if task_instance and task_instance.execution_date:
-        params['airflow.ctx.task_instance.execution_date'] = (
-            task_instance.execution_date.isoformat()
-        )
-
+        params[
+            AIRFLOW_VAR_NAME_FORMAT_MAPPING['AIRFLOW_CONTEXT_EXECUTION_DATE'][
+                name_format]] = task_instance.execution_date.isoformat()
+    dag_run = context.get('dag_run')
+    if dag_run and dag_run.run_id:
+        params[AIRFLOW_VAR_NAME_FORMAT_MAPPING['AIRFLOW_CONTEXT_DAG_RUN_ID'][
+            name_format]] = dag_run.run_id
     return params