[AIRFLOW-5390] Remove provide context

airflow/contrib/example_dags/example_qubole_operator.py

@@ -97,7 +97,6 @@ def compare_result(**kwargs):
 
 t3 = PythonOperator(
     task_id='compare_result',
-    provide_context=True,
     python_callable=compare_result,
     trigger_rule="all_done",
     dag=dag)

airflow/example_dags/docker_copy_data.py

@@ -69,7 +69,6 @@
 #
 # t_is_data_available = ShortCircuitOperator(
 #         task_id='check_if_data_available',
-#         provide_context=True,
 #         python_callable=is_data_available,
 #         dag=dag)
 #

airflow/example_dags/example_branch_python_dop_operator_3.py

@@ -58,7 +58,6 @@ def should_run(**kwargs):
 
 cond = BranchPythonOperator(
     task_id='condition',
-    provide_context=True,
     python_callable=should_run,
     dag=dag,
 )

airflow/example_dags/example_passing_params_via_test_command.py

@@ -37,17 +37,17 @@
 )
 
 
-def my_py_command(**kwargs):
+def my_py_command(test_mode, params):
     """
     Print out the "foo" param passed in via
     `airflow tasks test example_passing_params_via_test_command run_this <date>
     -tp '{"foo":"bar"}'`
     """
-    if kwargs["test_mode"]:
+    if test_mode:
         print(" 'foo' was passed in via test={} command : kwargs[params][foo] \
-               = {}".format(kwargs["test_mode"], kwargs["params"]["foo"]))
+               = {}".format(test_mode, params["foo"]))
     # Print out the value of "miff", passed in below via the Python Operator
-    print(" 'miff' was passed in via task params = {}".format(kwargs["params"]["miff"]))
+    print(" 'miff' was passed in via task params = {}".format(params["miff"]))
     return 1
 
 
@@ -58,7 +58,6 @@ def my_py_command(**kwargs):
 
 run_this = PythonOperator(
     task_id='run_this',
-    provide_context=True,
     python_callable=my_py_command,
     params={"miff": "agg"},
     dag=dag,

airflow/example_dags/example_python_operator.py

@@ -48,7 +48,6 @@ def print_context(ds, **kwargs):
 
 run_this = PythonOperator(
     task_id='print_the_context',
-    provide_context=True,
     python_callable=print_context,
     dag=dag,
 )

airflow/example_dags/example_trigger_target_dag.py

@@ -69,7 +69,6 @@ def run_this_func(**kwargs):
 
 run_this = PythonOperator(
     task_id='run_this',
-    provide_context=True,
     python_callable=run_this_func,
     dag=dag,
 )

airflow/example_dags/example_xcom.py

@@ -26,7 +26,6 @@
 args = {
     'owner': 'Airflow',
     'start_date': airflow.utils.dates.days_ago(2),
-    'provide_context': True,
 }
 
 dag = DAG('example_xcom', schedule_interval="@once", default_args=args)

airflow/operators/python_operator.py

@@ -25,6 +25,7 @@
 import types
 from textwrap import dedent
 from typing import Optional, Iterable, Dict, Callable
+from inspect import signature
 
 import dill
 
@@ -51,12 +52,6 @@ class PythonOperator(BaseOperator):
     :param op_args: a list of positional arguments that will get unpacked when
         calling your callable
     :type op_args: list (templated)
-    :param provide_context: if set to true, Airflow will pass a set of
-        keyword arguments that can be used in your function. This set of
-        kwargs correspond exactly to what you can use in your jinja
-        templates. For this to work, you need to define `**kwargs` in your
-        function header.
-    :type provide_context: bool
     :param templates_dict: a dictionary where the values are templates that
         will get templated by the Airflow engine sometime between
         ``__init__`` and ``execute`` takes place and are made available
@@ -79,7 +74,6 @@ def __init__(
         python_callable: Callable,
         op_args: Optional[Iterable] = None,
         op_kwargs: Optional[Dict] = None,
-        provide_context: bool = False,
         templates_dict: Optional[Dict] = None,
         templates_exts: Optional[Iterable[str]] = None,
         *args,
@@ -91,7 +85,6 @@ def __init__(
         self.python_callable = python_callable
         self.op_args = op_args or []
         self.op_kwargs = op_kwargs or {}
-        self.provide_context = provide_context
         self.templates_dict = templates_dict
         if templates_exts:
             self.template_ext = templates_exts
@@ -104,10 +97,21 @@ def execute(self, context):
                                  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
+        context.update(self.op_kwargs)
+        context['templates_dict'] = self.templates_dict
+
+        if {parameter for name, parameter
+           in signature(self.python_callable).parameters.items()
+           if str(parameter).startswith("**")}:
+            # If there is a **kwargs, **context or **_ then just pass everything.
             self.op_kwargs = context
+        else:
+            # If there is only for example, an execution_date, then pass only these in :-)
+            self.op_kwargs = {
+                name: context[name] for name, parameter
+                in signature(self.python_callable).parameters.items()
+                if name in context  # If it isn't available on the context, then ignore
+            }
 
         return_value = self.execute_callable()
         self.log.info("Done. Returned value was: %s", return_value)
@@ -130,6 +134,7 @@ class BranchPythonOperator(PythonOperator, SkipMixin):
     downstream to allow for the DAG state to fill up and the DAG run's state
     to be inferred.
     """
+
     def execute(self, context):
         branch = super().execute(context)
         self.skip_all_except(context['ti'], branch)
@@ -147,6 +152,7 @@ class ShortCircuitOperator(PythonOperator, SkipMixin):
 
     The condition is determined by the result of `python_callable`.
     """
+
     def execute(self, context):
         condition = super().execute(context)
         self.log.info("Condition result is %s", condition)
@@ -200,12 +206,6 @@ class PythonVirtualenvOperator(PythonOperator):
     :type op_kwargs: list
     :param op_kwargs: A dict of keyword arguments to pass to python_callable.
     :type op_kwargs: dict
-    :param provide_context: if set to true, Airflow will pass a set of
-        keyword arguments that can be used in your function. This set of
-        kwargs correspond exactly to what you can use in your jinja
-        templates. For this to work, you need to define `**kwargs` in your
-        function header.
-    :type provide_context: bool
     :param string_args: Strings that are present in the global var virtualenv_string_args,
         available to python_callable at runtime as a list[str]. Note that args are split
         by newline.
@@ -219,6 +219,7 @@ class PythonVirtualenvOperator(PythonOperator):
         processing templated fields, for examples ``['.sql', '.hql']``
     :type templates_exts: list[str]
     """
+
     @apply_defaults
     def __init__(
         self,
@@ -229,7 +230,6 @@ def __init__(
         system_site_packages: bool = True,
         op_args: Iterable = None,
         op_kwargs: Dict = None,
-        provide_context: bool = False,
         string_args: Optional[Iterable[str]] = None,
         templates_dict: Optional[Dict] = None,
         templates_exts: Optional[Iterable[str]] = None,
@@ -242,7 +242,6 @@ def __init__(
             op_kwargs=op_kwargs,
             templates_dict=templates_dict,
             templates_exts=templates_exts,
-            provide_context=provide_context,
             *args,
             **kwargs)
         self.requirements = requirements or []
@@ -383,7 +382,7 @@ def _generate_python_code(self):
         fn = self.python_callable
         # dont try to read pickle if we didnt pass anything
         if self._pass_op_args():
-            load_args_line = 'with open(sys.argv[1], "rb") as file: arg_dict = {}.load(file)'\
+            load_args_line = 'with open(sys.argv[1], "rb") as file: arg_dict = {}.load(file)' \
                 .format(pickling_library)
         else:
             load_args_line = 'arg_dict = {"args": [], "kwargs": {}}'

docs/howto/operator/python.rst

@@ -44,9 +44,9 @@ to the Python callable.
 Templating
 ^^^^^^^^^^
 
-When you set the ``provide_context`` argument to ``True``, Airflow passes in
-an additional set of keyword arguments: one for each of the :doc:`Jinja
-template variables <../../macros-ref>` and a ``templates_dict`` argument.
+Airflow passes in an additional set of keyword arguments: one for each of the
+:doc:`Jinja template variables <../../macros-ref>` and a ``templates_dict``
+argument.
 
 The ``templates_dict`` argument is templated, so each value in the dictionary
 is evaluated as a :ref:`Jinja template <jinja-templating>`.

tests/contrib/operators/test_aws_athena_operator.py

@@ -53,7 +53,6 @@ def setUp(self):
         args = {
             'owner': 'airflow',
             'start_date': DEFAULT_DATE,
-            'provide_context': True
         }
 
         self.dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once',

tests/contrib/operators/test_s3_to_sftp_operator.py

@@ -69,7 +69,6 @@ def setUp(self):
         args = {
             'owner': 'airflow',
             'start_date': DEFAULT_DATE,
-            'provide_context': True
         }
         dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args)
         dag.schedule_interval = '@once'

tests/contrib/operators/test_sftp_operator.py

@@ -53,7 +53,6 @@ def setUp(self):
         args = {
             'owner': 'airflow',
             'start_date': DEFAULT_DATE,
-            'provide_context': True
         }
         dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args)
         dag.schedule_interval = '@once'

tests/contrib/operators/test_sftp_to_s3_operator.py

@@ -68,7 +68,6 @@ def setUp(self):
         args = {
             'owner': 'airflow',
             'start_date': DEFAULT_DATE,
-            'provide_context': True
         }
         dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args)
         dag.schedule_interval = '@once'

tests/contrib/operators/test_ssh_operator.py

@@ -51,7 +51,6 @@ def setUp(self):
         args = {
             'owner': 'airflow',
             'start_date': DEFAULT_DATE,
-            'provide_context': True
         }
         dag = DAG(TEST_DAG_ID + 'test_schedule_dag_once', default_args=args)
         dag.schedule_interval = '@once'

tests/core.py

@@ -568,7 +568,6 @@ def test_py_op(templates_dict, ds, **kwargs):
 
         t = PythonOperator(
             task_id='test_py_op',
-            provide_context=True,
             python_callable=test_py_op,
             templates_dict={'ds': "{{ ds }}"},
             dag=self.dag)

tests/dags/test_cli_triggered_dags.py

@@ -50,6 +50,5 @@ def success(ti=None, *args, **kwargs):
 dag1_task2 = PythonOperator(
     task_id='test_run_dependent_task',
     python_callable=success,
-    provide_context=True,
     dag=dag1)
 dag1_task1.set_downstream(dag1_task2)

tests/operators/test_python_operator.py

@@ -57,8 +57,8 @@ def build_recording_function(calls_collection):
     Then using this custom function recording custom Call objects for further testing
     (replacing Mock.assert_called_with assertion method)
     """
-    def recording_function(*args, **kwargs):
-        calls_collection.append(Call(*args, **kwargs))
+    def recording_function(*args):
+        calls_collection.append(Call(*args))
     return recording_function
 
 
@@ -129,11 +129,10 @@ def test_python_operator_python_callable_is_callable(self):
                 task_id='python_operator',
                 dag=self.dag)
 
-    def _assertCallsEqual(self, first, second):
+    def _assert_calls_equal(self, first, second):
         self.assertIsInstance(first, Call)
         self.assertIsInstance(second, Call)
         self.assertTupleEqual(first.args, second.args)
-        self.assertDictEqual(first.kwargs, second.kwargs)
 
     def test_python_callable_arguments_are_templatized(self):
         """Test PythonOperator op_args are templatized"""
@@ -148,7 +147,7 @@ def test_python_callable_arguments_are_templatized(self):
             task_id='python_operator',
             # a Mock instance cannot be used as a callable function or test fails with a
             # TypeError: Object of type Mock is not JSON serializable
-            python_callable=(build_recording_function(recorded_calls)),
+            python_callable=build_recording_function(recorded_calls),
             op_args=[
                 4,
                 date(2019, 1, 1),
@@ -167,7 +166,7 @@ def test_python_callable_arguments_are_templatized(self):
 
         ds_templated = DEFAULT_DATE.date().isoformat()
         self.assertEqual(1, len(recorded_calls))
-        self._assertCallsEqual(
+        self._assert_calls_equal(
             recorded_calls[0],
             Call(4,
                  date(2019, 1, 1),
@@ -183,7 +182,7 @@ def test_python_callable_keyword_arguments_are_templatized(self):
             task_id='python_operator',
             # a Mock instance cannot be used as a callable function or test fails with a
             # TypeError: Object of type Mock is not JSON serializable
-            python_callable=(build_recording_function(recorded_calls)),
+            python_callable=build_recording_function(recorded_calls),
             op_kwargs={
                 'an_int': 4,
                 'a_date': date(2019, 1, 1),
@@ -200,7 +199,7 @@ def test_python_callable_keyword_arguments_are_templatized(self):
         task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
 
         self.assertEqual(1, len(recorded_calls))
-        self._assertCallsEqual(
+        self._assert_calls_equal(
             recorded_calls[0],
             Call(an_int=4,
                  a_date=date(2019, 1, 1),

tests/operators/test_virtualenv_operator.py

@@ -199,20 +199,6 @@ def f(_):
         self._run_as_operator(f, op_args=[datetime.datetime.utcnow()])
 
     def test_context(self):
-        def f(**kwargs):
-            return kwargs['templates_dict']['ds']
+        def f(templates_dict):
+            return templates_dict['ds']
         self._run_as_operator(f, templates_dict={'ds': '{{ ds }}'})
-
-    def test_provide_context(self):
-        def fn():
-            pass
-        task = PythonVirtualenvOperator(
-            python_callable=fn,
-            python_version=sys.version_info[0],
-            task_id='task',
-            dag=self.dag,
-            provide_context=True,
-        )
-        self.assertTrue(
-            task.provide_context
-        )

tests/utils/test_log_handlers.py

@@ -71,7 +71,6 @@ def task_callable(ti, **kwargs):
             task_id='task_for_testing_file_log_handler',
             dag=dag,
             python_callable=task_callable,
-            provide_context=True
         )
         ti = TaskInstance(task=task, execution_date=DEFAULT_DATE)
 
@@ -123,7 +122,6 @@ def task_callable(ti, **kwargs):
             task_id='task_for_testing_file_log_handler',
             dag=dag,
             python_callable=task_callable,
-            provide_context=True
         )
         ti = TaskInstance(task=task, execution_date=DEFAULT_DATE)
         ti.try_number = 2