`context` doesn't contain `TaskIntance` when used in `on_*_callback` at the `DAG` level

[jnstockley]

Apache Airflow version

3.0.1

If "Other Airflow 2 version" selected, which one?

No response

What happened?

When using the on_*_callback at the dag level, the context that gets passed to the notify function only contains these keys

• dag
• run_id
• reason

When it's called at a task level these keys are present

• dag
• inlets
• map_index_template
• outlets
• run_id
• task
• task_instance
• ti
• outlet_events
• inlet_events
• macros
• params
• var
• conn
• dag_run
• triggering_asset_events
• task_instance_key_str
• task_reschedule_count
• prev_start_date_success
• prev_end_date_success
• logical_date
• ds
• ds_nodash
• ts
• ts_nodash
• ts_nodash_with_tz
• data_interval_end
• data_interval_start
• prev_data_interval_start_success
• prev_data_interval_end_success
• templates_dict

What you think should happen instead?

The keys should be the same when used at a DAG level and a Task level, and should populated similarly, i.e., for task related keys should use the last task that completed in a DAG

How to reproduce

Create a custom notifier to be able to see the contents of context. Use the on_*_callback parameter at the DAG at task level using the TaskFlowAPI, not sure if this matters or not.

Operating System

Ubuntu

Versions of Apache Airflow Providers

No response

Deployment

Other Docker-based deployment

Deployment details

Custom docker container using apache/airflow:3.0.2rc1 as the base

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[jroachgolf84]

@jnstockley, can you provide a snippet of your DAG-definition? I'm mostly curious in which callbacks you are setting at the DAG-level.

[jnstockley]

@jroachgolf84 I tired it with on_success_callback and on_failure_callback and got the same results.
`

import json

import pendulum

from airflow.sdk import dag, task

from plugins.notifications.custom import CustomNotifier


@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example"],
    on_success_callback=CustomNotifier(to="email-address"),
)
def tutorial_taskflow_api():`

[jroachgolf84]

This is the expected behavior - on_success_callback and on_failure_callback are applied when a DAG enters a "Success" or "Failed" terminal state, not when a Task does. This means that no TaskInstance-level information would be available here.

The docs are a bit mis-leading, and I've got an open PR to address this: https://github.com/apache/airflow/pull/51150/files

[jroachgolf84]

That being said, I think there are items on that list that you shared that would be worth thinking about bringing over...

[jnstockley]

That makes sense but it also seems like something that changed between V2 and V3. Is there a way to get the V2 behavior in V3, i.e., a different parameter to use at the DAG level?

[jroachgolf84]

@jnstockley, this is where this context is set.

https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/models/dagrun.py#L1327-L1333

It appears that this change was introduced last month by this PR: https://github.com/apache/airflow/pull/50300/files.

[jroachgolf84]

Are there additional fields that you are interested in bring in? Hypothetically, this would be a very light-lift to implement.

[jnstockley]

@jroachgolf84 Ideally I would love either task_instance or ti or at least these fields from task_instance

• state I'm assuming reason would be the same but I could be wrong
• try_number
• max_tries

[jroachgolf84]

I think the hard thing here is that we don't have information about the Task that is causing this DAG to move into a terminal state (whether that is success or failure) in the object itself. Was this functionality something that you had working in a certain version of Airflow 2.X?