Repository: apache/airflow
Pull request: Add function to get current context
Support for getting current context at any codelocation that runs under the scope of BaseOperator.execute function. This functionality is part of AIP-31

closes: #8058 

Work based on @jonathanshir PR #8651 

---
Make sure to mark the boxes below before creating PR: [x]

- [x] Description above provides context of the change
- [x] Unit tests coverage for changes (not needed for documentation changes)
- [x] Target Github ISSUE in description if exists
- [x] Commits follow "[How to write a good git commit message](http://chris.beams.io/posts/git-commit/)"
- [x] Relevant documentation is updated including usage instructions.
- [x] I will engage committers as explained in [Contribution Workflow Example](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#contribution-workflow-example).

---
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).
Read the [Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines) for more information.


Problem statement:
[AIP-31] Add lazy way to get Airflow context from a wrapped function
**Description**

Currently the only way to access the context object from a PythonOperator wrapped function is by setting `provide_context=True` in the Operator. That loads it into the kwargs. This makes the signature of the function become difficult to argue about.

Adding a way to lazy get the function when executing the function may be better. Similar to Flask `flask.request` that get populated when the function is executed.

**Use case / motivation**

Use the Airflow context in arbitrary function while keeping the signature of the function stable and easy to reason about. 

An idea of implementation would be:

```
from airflow.task import context

@task 
def context_task(...)
    return context['run_date']
```

This can be populated by enabling a context manager on the `PythonOperator.execute` function that populates the field (similar to how `DagContext` works) and removes it afterwards.

Hints from issue discussion:
Just a comment, provide_context is no longer requires in Master: https://github.com/apache/airflow/pull/5990/files/a5d99ed3f3c7b1a815bd2617d3b935bed499f3f7
I see, I still think we should implement this lazy resolved context. Using the current approach from master means that signatures can mismatch op_args/op_kwargs and would make it difficult to reason about the function signature:

```
@task 
def simple_task(text:str, execution_date: datetime):
  print(text)
  print(execution_date)

simple_task("hello")
```
As you can see this will execute correctly even though the function signature has a an extra argument. This leads to confusion in my personal opinion hence my proposed solution.


> I see, I still think we should implement this lazy resolved context. Using the current approach from master means that signatures can mismatch op_args/op_kwargs and would make it difficult to reason about the function signature:
> 
> ```
> @task 
> def simple_task(text:str, execution_date: datetime):
>   print(text)
>   print(execution_date)
> 
> simple_task("hello")
> ```
> 
> As you can see this will execute correctly even though the function signature has a an extra argument. This leads to confusion in my personal opinion hence my proposed solution.

Yup agree, just wanted to make it clear that we don't need provide_context anymore
My suggested changes:
The issue states we want to implement context similarly to flask.requests :
```
from airflow.task import context
@task 
def context_task(...)
    return context['run_date']
```
I think this is a bit of an over-complication - our use case is much simpler than this.
We simply have a single execute context for each task instance, so my intended implementation will be a bit simpler:
```
from airflow.current import current_context
@task
def context_task(...)
    context = current_context()
    return context['run_date']
```
The actual implementation will contain a simple stack data type that is updated before op.execute is called.
Thoughts?
> We simply have a single execute context for each task instance, so my intended implementation will be a bit simpler:
> 
> ```
> from airflow.current import current_context
> @task
> def context_task(...)
>     context = current_context()
>     return context['run_date']
> ```

+1 for starting simple, what would you say to make this context a namedtuple? Referencing values is much more easier with namedtuple than with dictionary (no need to remember names).
Sounds great, but does this imply we also change the context being provided the old way?
```
@provide_context
def task(**context):
    print(type(context)) # == NamedTuple?
```
> Sounds great, but does this imply we also change the context being provided the old way?

No, if someone uses `**context` then it implicitly suggest that it will be a `dict`. I would do this only for the "new" context. Also, we can alter namedtuple  `__getitem__` to support old (dict like) behaviour.


Saw it on the PR, lgtm as well on making it a function 😄 
Any updates on this?

Commit message: Add function to get current context

Support for getting current context at any codelocation that runs
under the scope of BaseOperator.execute function. This functionality
is part of AIP-31

Patch to review:
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 8360097e0ec5a..0ded0ea1fd0e0 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -51,6 +51,7 @@
 from airflow.sentry import Sentry
 from airflow.settings import STORE_SERIALIZED_DAGS
 from airflow.stats import Stats
+from airflow.task.context.current import set_current_context
 from airflow.ti_deps.dep_context import DepContext
 from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS
 from airflow.utils import timezone
@@ -1132,7 +1133,8 @@ def signal_handler(signum, frame):  # pylint: disable=unused-argument
         self._run_execute_callback(context, task)
 
         # Execute the task
-        result = self._execute_task(context, task_copy)
+        with set_current_context(context):
+            result = self._execute_task(context, task_copy)
 
         # Run post_execute callback
         task_copy.post_execute(context=context, result=result)
diff --git a/airflow/task/context/__init__.py b/airflow/task/context/__init__.py
new file mode 100644
index 0000000000000..13a83393a9124
--- /dev/null
+++ b/airflow/task/context/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git a/airflow/task/context/current.py b/airflow/task/context/current.py
new file mode 100644
index 0000000000000..d5db2e8ecb25e
--- /dev/null
+++ b/airflow/task/context/current.py
@@ -0,0 +1,69 @@
+# 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 contextlib
+import logging
+from typing import Any, Dict
+
+from airflow.exceptions import AirflowException
+
+_CURRENT_CONTEXT = []
+log = logging.getLogger(__name__)
+
+
+@contextlib.contextmanager
+def set_current_context(context: Dict[str, Any]):
+    """
+    Sets the current execution context to the provided context object.
+    This method should be called once per Task execution, before calling operator.execute
+    """
+    _CURRENT_CONTEXT.append(context)
+    try:
+        yield context
+    finally:
+        expected_state = _CURRENT_CONTEXT.pop()
+        if expected_state != context:
+            log.warning(
+                "Current context is not equal to the state at context stack. Expected=%s, got=%s",
+                context,
+                expected_state,
+            )
+
+
+def get_current_context() -> Dict[str, Any]:
+    """
+    Obtain the execution context for the currently executing operator without
+    altering user method's signature.
+    This is the simplest method of retrieving the execution context dictionary.
+    ** Old style:
+        def my_task(**context):
+            ti = context["ti"]
+    ** New style:
+        from airflow.task.context import get_current_context
+        def my_task():
+            context = get_current_context()
+            ti = context["ti"]
+
+    Current context will only have value if this method was called after an operator
+    was starting to execute.
+    """
+    if not _CURRENT_CONTEXT:
+        raise AirflowException(
+            "Current context was requested but no context was found! "
+            "Are you running within an airflow task?"
+        )
+    return _CURRENT_CONTEXT[-1]
diff --git a/tests/task/context/__init__.py b/tests/task/context/__init__.py
new file mode 100644
index 0000000000000..13a83393a9124
--- /dev/null
+++ b/tests/task/context/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git a/tests/task/context/test_current_context.py b/tests/task/context/test_current_context.py
new file mode 100644
index 0000000000000..fdb715f046244
--- /dev/null
+++ b/tests/task/context/test_current_context.py
@@ -0,0 +1,101 @@
+# 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.
+
+from datetime import datetime, timedelta
+
+import pytest
+
+from airflow import DAG
+from airflow.exceptions import AirflowException
+from airflow.models.baseoperator import BaseOperator
+from airflow.operators.python import PythonOperator
+from airflow.task.context.current import get_current_context, set_current_context
+from airflow.utils.dates import days_ago
+
+DEFAULT_ARGS = {
+    "owner": "test",
+    "depends_on_past": True,
+    "start_date": days_ago(1),
+    "end_date": datetime.today(),
+    "schedule_interval": "@once",
+    "retries": 1,
+    "retry_delay": timedelta(minutes=1),
+}
+
+
+class TestCurrentContext:
+    def test_current_context_no_context_raise(self):
+        with pytest.raises(AirflowException):
+            get_current_context()
+
+    def test_current_context_roundtrip(self):
+        example_context = {"Hello": "World"}
+
+        with set_current_context(example_context):
+            assert get_current_context() == example_context
+
+    def test_context_removed_after_exit(self):
+        example_context = {"Hello": "World"}
+
+        with set_current_context(example_context):
+            pass
+        with pytest.raises(AirflowException, ):
+            get_current_context()
+
+    def test_nested_context(self):
+        """
+        Nested execution context should be supported in case the user uses multiple context managers.
+        Each time the execute method of an operator is called, we set a new 'current' context.
+        This test verifies that no matter how many contexts are entered - order is preserved
+        """
+        max_stack_depth = 15
+        ctx_list = []
+        for i in range(max_stack_depth):
+            # Create all contexts in ascending order
+            new_context = {"ContextId": i}
+            # Like 15 nested with statements
+            ctx_obj = set_current_context(new_context)
+            ctx_obj.__enter__()  # pylint: disable=E1101
+            ctx_list.append(ctx_obj)
+        for i in reversed(range(max_stack_depth)):
+            # Iterate over contexts in reverse order - stack is LIFO
+            ctx = get_current_context()
+            assert ctx["ContextId"] == i
+            # End of with statement
+            ctx_list[i].__exit__(None, None, None)
+
+
+class MyContextAssertOperator(BaseOperator):
+    def execute(self, context):
+        assert context == get_current_context()
+
+
+def get_all_the_context(**context):
+    current_context = get_current_context()
+    assert context == current_context
+
+
+class TestCurrentContextRuntime:
+    def test_context_in_task(self):
+        with DAG(dag_id="assert_context_dag", default_args=DEFAULT_ARGS):
+            op = MyContextAssertOperator(task_id="assert_context")
+            op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)
+
+    def test_get_context_in_old_style_context_task(self):
+        with DAG(dag_id="edge_case_context_dag", default_args=DEFAULT_ARGS):
+            op = PythonOperator(python_callable=get_all_the_context, task_id="get_all_the_context")
+            op.run(ignore_first_depends_on_past=True, ignore_ti_state=True)


Code review comments:
diff_hunk: @@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
path: airflow/task/context/current.py
text: @user1:
Hmm.. I feel this is wrong place for this functions. I agree "utils" isn't the best place but mostly things under "airflow/task" are for task_runners but  I don't know where we can put it.🤷‍♂️

@author:
How about `operator/python.py`?  This will make import simple:
```python
from airflow.operators.python import task, get_current_context
```