Skip to content

Conversation

@kaxil
Copy link
Member

@kaxil kaxil commented Jan 8, 2025

closes #45234

I am putting the logic for set_current_context in execution_time/context.py. I didn't want to put _CURRENT_CONTEXT in task_sdk/src/airflow/sdk/definitions/contextmanager.py to avoid execution logic in a user-facing module but I couldn't think of another way to store it from execution & allow retrieving (via get_current_context in the Standard Provider) in their Task.

Upcoming PRs:

  • Move most of the internal stuff in Task SDK to a separate module.
  • Use create_runtime_ti fixture more widely in tests

Tested with the following DAG:

import pendulum

from airflow.decorators import dag, task
from airflow.providers.standard.operators.python import get_current_context

@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
)
def x_get_context():
    @task
    def template_test(data_interval_end):
        context = get_current_context()

        # Will print `2024-10-10 00:00:00+00:00`.
        # Note how we didn't pass this value when calling the task. Instead
        # it was passed by the decorator from the context
        print(f"data_interval_end: {data_interval_end}")

        # Will print the full context dict
        print(f"context: {context}")

    template_test()

x_get_context()
image

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

@kaxil kaxil added area:task-execution-interface-aip72 AIP-72: Task Execution Interface (TEI) aka Task SDK and removed area:providers provider:standard labels Jan 8, 2025
@kaxil kaxil force-pushed the get-current-context branch 2 times, most recently from 06b7acd to 66526e0 Compare January 8, 2025 20:10
@kaxil kaxil marked this pull request as ready for review January 8, 2025 20:10
@kaxil kaxil requested review from XD-DENG and ashb as code owners January 8, 2025 20:10
@kaxil kaxil force-pushed the get-current-context branch from 66526e0 to 9627d5a Compare January 8, 2025 20:14
@kaxil kaxil requested a review from amoghrajesh January 8, 2025 20:16
@kaxil kaxil force-pushed the get-current-context branch 2 times, most recently from 9f96a41 to 407002b Compare January 9, 2025 06:04
closes apache#45234

I am putting the logic for `set_current_context` in `execution_time/context.py`. I didn't want to put `_CURRENT_CONTEXT` in `task_sdk/src/airflow/sdk/definitions/contextmanager.py` to avoid execution logic in a user-facing module but I couldn't think of another way to store it from execution & allow retrieving (via `get_current_context` in the Standard Provider) in their Task.

Upcoming PRs:
- Move most of the internal stuff in Task SDK to a separate module.
- Use `create_runtime_ti` fixture more widely in tests

---

Tested with the following DAG:

```py
import pendulum

from airflow.decorators import dag, task
from airflow.providers.standard.operators.python import get_current_context

@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
)
def x_get_context():
    @task
    def template_test(data_interval_end):
        context = get_current_context()

        # Will print `2024-10-10 00:00:00+00:00`.
        # Note how we didn't pass this value when calling the task. Instead
        # it was passed by the decorator from the context
        print(f"data_interval_end: {data_interval_end}")

        # Will print the full context dict
        print(f"context: {context}")

    template_test()

x_get_context()

```
<img width="1703" alt="image" src="https://github.com/user-attachments/assets/2763963a-d299-412f-bee3-3b20904ca7c8" />
@kaxil kaxil force-pushed the get-current-context branch from 407002b to 14de0a6 Compare January 9, 2025 06:43
@kaxil kaxil merged commit 0480623 into apache:main Jan 9, 2025
63 checks passed
@kaxil kaxil deleted the get-current-context branch January 9, 2025 07:15
kaxil added a commit to astronomer/airflow that referenced this pull request Jan 9, 2025
I had added `create_runtime_ti` in apache#45486, this PR modifies the fixture a little bit and uses it more broadly.
kaxil added a commit to astronomer/airflow that referenced this pull request Jan 9, 2025
I had added `create_runtime_ti` in apache#45486, this PR modifies the fixture a little bit and uses it more broadly.
kaxil added a commit that referenced this pull request Jan 9, 2025
I had added `create_runtime_ti` in #45486, this PR modifies the fixture a little bit and uses it more broadly.
Comment on lines +101 to +113
This fixture sets up a `RuntimeTaskInstance` with default or custom `TIRunContext` and `StartupDetails`,
making it easy to simulate task execution scenarios in tests.

Example usage: ::

def test_custom_task_instance(create_runtime_ti):
class MyTaskOperator(BaseOperator):
def execute(self, context):
assert context["dag_run"].run_id == "test_run"

task = MyTaskOperator(task_id="test_task")
ti = create_runtime_ti(task, context_from_server=make_ti_context(run_id="test_run"))
# Further test logic...
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Amazing!

karenbraganz pushed a commit to karenbraganz/airflow that referenced this pull request Jan 13, 2025
closes apache#45234

I am putting the logic for `set_current_context` in `execution_time/context.py`. I didn't want to put `_CURRENT_CONTEXT` in `task_sdk/src/airflow/sdk/definitions/contextmanager.py` to avoid execution logic in a user-facing module but I couldn't think of another way to store it from execution & allow retrieving (via `get_current_context` in the Standard Provider) in their Task.

Upcoming PRs:
- Move most of the internal stuff in Task SDK to a separate module.
- Use `create_runtime_ti` fixture more widely in tests

---

Tested with the following DAG:

```py
import pendulum

from airflow.decorators import dag, task
from airflow.providers.standard.operators.python import get_current_context

@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
)
def x_get_context():
    @task
    def template_test(data_interval_end):
        context = get_current_context()

        # Will print `2024-10-10 00:00:00+00:00`.
        # Note how we didn't pass this value when calling the task. Instead
        # it was passed by the decorator from the context
        print(f"data_interval_end: {data_interval_end}")

        # Will print the full context dict
        print(f"context: {context}")

    template_test()

x_get_context()

```
<img width="1703" alt="image" src="https://github.com/user-attachments/assets/2763963a-d299-412f-bee3-3b20904ca7c8" />
karenbraganz pushed a commit to karenbraganz/airflow that referenced this pull request Jan 13, 2025
I had added `create_runtime_ti` in apache#45486, this PR modifies the fixture a little bit and uses it more broadly.
HariGS-DB pushed a commit to HariGS-DB/airflow that referenced this pull request Jan 16, 2025
closes apache#45234

I am putting the logic for `set_current_context` in `execution_time/context.py`. I didn't want to put `_CURRENT_CONTEXT` in `task_sdk/src/airflow/sdk/definitions/contextmanager.py` to avoid execution logic in a user-facing module but I couldn't think of another way to store it from execution & allow retrieving (via `get_current_context` in the Standard Provider) in their Task.

Upcoming PRs:
- Move most of the internal stuff in Task SDK to a separate module.
- Use `create_runtime_ti` fixture more widely in tests

---

Tested with the following DAG:

```py
import pendulum

from airflow.decorators import dag, task
from airflow.providers.standard.operators.python import get_current_context

@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
)
def x_get_context():
    @task
    def template_test(data_interval_end):
        context = get_current_context()

        # Will print `2024-10-10 00:00:00+00:00`.
        # Note how we didn't pass this value when calling the task. Instead
        # it was passed by the decorator from the context
        print(f"data_interval_end: {data_interval_end}")

        # Will print the full context dict
        print(f"context: {context}")

    template_test()

x_get_context()

```
<img width="1703" alt="image" src="https://github.com/user-attachments/assets/2763963a-d299-412f-bee3-3b20904ca7c8" />
HariGS-DB pushed a commit to HariGS-DB/airflow that referenced this pull request Jan 16, 2025
I had added `create_runtime_ti` in apache#45486, this PR modifies the fixture a little bit and uses it more broadly.
got686-yandex pushed a commit to got686-yandex/airflow that referenced this pull request Jan 30, 2025
closes apache#45234

I am putting the logic for `set_current_context` in `execution_time/context.py`. I didn't want to put `_CURRENT_CONTEXT` in `task_sdk/src/airflow/sdk/definitions/contextmanager.py` to avoid execution logic in a user-facing module but I couldn't think of another way to store it from execution & allow retrieving (via `get_current_context` in the Standard Provider) in their Task.

Upcoming PRs:
- Move most of the internal stuff in Task SDK to a separate module.
- Use `create_runtime_ti` fixture more widely in tests

---

Tested with the following DAG:

```py
import pendulum

from airflow.decorators import dag, task
from airflow.providers.standard.operators.python import get_current_context

@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
)
def x_get_context():
    @task
    def template_test(data_interval_end):
        context = get_current_context()

        # Will print `2024-10-10 00:00:00+00:00`.
        # Note how we didn't pass this value when calling the task. Instead
        # it was passed by the decorator from the context
        print(f"data_interval_end: {data_interval_end}")

        # Will print the full context dict
        print(f"context: {context}")

    template_test()

x_get_context()

```
<img width="1703" alt="image" src="https://github.com/user-attachments/assets/2763963a-d299-412f-bee3-3b20904ca7c8" />
got686-yandex pushed a commit to got686-yandex/airflow that referenced this pull request Jan 30, 2025
I had added `create_runtime_ti` in apache#45486, this PR modifies the fixture a little bit and uses it more broadly.
kosteev pushed a commit to GoogleCloudPlatform/composer-airflow that referenced this pull request May 28, 2025
I had added `create_runtime_ti` in apache/airflow#45486, this PR modifies the fixture a little bit and uses it more broadly.

GitOrigin-RevId: 0e114c2a70642c17b0be7ff76c5d86f4145253b7
kosteev pushed a commit to GoogleCloudPlatform/composer-airflow that referenced this pull request Sep 23, 2025
I had added `create_runtime_ti` in apache/airflow#45486, this PR modifies the fixture a little bit and uses it more broadly.

GitOrigin-RevId: 0e114c2a70642c17b0be7ff76c5d86f4145253b7
kosteev pushed a commit to GoogleCloudPlatform/composer-airflow that referenced this pull request Oct 21, 2025
I had added `create_runtime_ti` in apache/airflow#45486, this PR modifies the fixture a little bit and uses it more broadly.

GitOrigin-RevId: 0e114c2a70642c17b0be7ff76c5d86f4145253b7
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:task-execution-interface-aip72 AIP-72: Task Execution Interface (TEI) aka Task SDK area:task-sdk

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add support for get_current_context in Task SDK

3 participants