Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

AIP-72: Add support for get_current_context in Task SDK #45486

Merged
merged 1 commit into from
Jan 9, 2025

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!

agupta01 pushed a commit to agupta01/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" />
agupta01 pushed a commit to agupta01/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.
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.
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