Pull the extra argument from Dataset

[Vosec]

Hello everyone,

I was wondering If I could somehow pull the extra argument from the defined Dataset.

If I understand this concept correctly, the Dataset object is saved in the airflow backend DB along with the extra parameter. For example:

ds1 = Dataset("/tmp/airflow/file1.csv", extra={"team": "trainees"})
The same Dataset is defined in the consumer DAG, but without the extra argument.

Is there any way to pull the extra from producer DAG in the consumer DAG? Or the extra argument is linked with the specific DAG and there is no way of pulling it?
If not, I can make workaround to put some metadata in the specified URI...

Thank you for your time.

[sgoncharov1]

Hi,
We also would be happy for such functionality.
Scenario: we have a Dag that listens for an event of data availability for specific date/hour.
Once the data is available, this Dag notifies a change on certain dataset (and we would like to add date/hour as extra parameters).
On the consumer side, we have other Dags that listen for these events, and when they are triggered , they want to know for which hour. This extra info also enables doing reprocessing of past runs.

[Blizzke]

We are experiencing the same issue and have been looking int it.
Unfortunately it all goes a bit deeper than just 'retrieving the extra'.

Airflow uses the Dataset instance (dataset table) for it's id everywhere.
So, if you happen to create a new dataset that has extra defined, it will be stored in that table as well, but afterwards it really doesn't matter anymore.

For us, that isn't really important, as that is how we want to use it (our extra can differ per task instance that produces an outlet, but the URI is what binds everything together).
Seems they actually thought about that, by giving the DatasetEvent-model (dataset_event table) an extra field as well.
Unfortunately, that is never used.

TaskInstance._run_raw_task() calls its _register_dataset_changes() function at the end and passes the outlet instances as is, without specifiying a separate extra (which the DatasetManager requires if you want to store it in the event). The extra that is part of the outlet/dataset object, is never used, only the uri is.
So unless it was a completely new dataset, you've already lost the data right there.

This is easily resolved with a custom DatasetManager (the original one is airflow.datasets.manager.DatasetManager).
The only thing we need is to override the register_dataset_change function and add 2 lines to it:

class DatasetManager(OriginalDatasetManager):
    def register_dataset_change(self, *, task_instance: TaskInstance, dataset: Dataset, extra=None, session: Session,
                                **kwargs) -> None:
        if extra is None and dataset.extra:
            extra = dataset.extra

        super().register_dataset_change(task_instance=task_instance, dataset=dataset, extra=extra, session=session, **kwargs)

(It's pretty easy to override this one in the conf file, add dataset_manager_class = whatever_package.DatasetManager to the [core]-section)

By the way, yes, it will essentially duplicate the extra for new datasets between the dataset and the dataset_event table, but we don't really care about that.

Now we're halfway, we still need to retrieve the extra during a run.

Whenever a new dag run is triggered by dataset events, it will receive a set of DatasetEvent-instances that caused it to trigger via way of consumed_dataset_events in the DagRun instance (courtesy of SchedulerJobRunner._create_dag_runs_dataset_triggered).
Just fetching it causes a DetachedInstanceError and really didn't want to mess with SQLAlchemy sessions.
So after more searching I found out these are also added to the templating context.
With that info I wrote a helper function that abuses that to obtain it and convert into a set of airflow.datasets.Dataset instances:

def get_triggering_datasets(context: Context = None) -> List[Dataset]:
    from airflow.models.dataset import DatasetEvent
    from airflow.datasets import Dataset

    if not context:
        from airflow.operators.python import get_current_context
        context = get_current_context()

    uri: str
    event: DatasetEvent
    datasets: List[Dataset] = []

    task_instance = context['task_instance'].get_template_context()["triggering_dataset_events"]
    for uri, events in task_instance.items():
        if not isinstance(events, list):
            continue

        for event in events:
            if not isinstance(event, DatasetEvent):
                continue

            ds = Dataset(uri=uri, extra=event.extra or None)
            # Avoid duplicates (it is possible we get multiple events that contain the same uri + extra)
            if ds not in datasets:
                datasets.append(ds)

    return datasets

What does this do?
Within the TaskInstance templating context, there is a variable triggering_dataset_events, as you can see.
This list is keyed by uri, but has other content, like a default_factory, hence the extra instanceof checks to skip these.
Every uri-key has a list of DatasetEvent instances.
All combined we just return the internally used Dataset instances.

Since it uses the task instance, you can only use it within operators, but that's all we need.
Data wise, you're not bound to the task instance, the actual events are linked to the DagRun.
Just look at the get_triggering_events function (which is a subfunction of TaskInstance.get_template_context() on how to do it.

Sorry for the long write-up (doing this as much for my own memory as anything else), but have fun with it.

[sgoncharov1]

@Blizzke thank you very much for your detailed advise.
We will try this approach!

[ketozhang]

Looks like a more concise example has been added to the docs.

https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/datasets.html#fetching-information-from-a-triggering-dataset-event

[RNHTTR]

I believe #37006 will fix this, right @Blizzke ?

[fritz-astronomer]

Not sure I have a better answer, but for potentially similar reasons I am using XCOMs from the originating task to pass data
https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/datasets.html#fetching-information-from-a-triggering-dataset-event