Skip to content

Add support for get_current_context in Task SDK #45234

@kaxil

Description

@kaxil

The following task should work. This task is part of tutorial_taskflow_templates DAG

    from airflow.providers.standard.operators.python import get_current_context

    @task(
        # Causes variables that end with `.sql` to be read and templates
        # within to be rendered.
        templates_exts=[".sql"],
    )
    def template_test(sql, test_var, data_interval_end):
        context = get_current_context()

        # Will print...
        # select * from test_data
        # where 1=1
        #     and run_id = 'scheduled__2024-10-09T00:00:00+00:00'
        #     and something_else = 'param_from_task'
        print(f"sql: {sql}")

        # Will print `scheduled__2024-10-09T00:00:00+00:00`
        print(f"test_var: {test_var}")

        # 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...
        # run_id: scheduled__2024-10-09T00:00:00+00:00; params.other_param: from_dag
        template_str = "run_id: {{ run_id }}; params.other_param: {{ params.other_param }}"
        rendered_template = context["task"].render_template(
            template_str,
            context,
        )
        print(f"rendered template: {rendered_template}")

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

Metadata

Metadata

Assignees

Type

No type

Projects

Status

Done

Milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions