AIP-66 Refactor DagRun to DagVersion association

[ephraimbuddy]

This removes DagRun.dag_version association and replaces it with dag_versions property on DagRun that collects all the dag_version_ids associated with the task instances of the DagRun.

closes: #46565

[jedcunningham]

We should probably union this instead so it's a single query vs 2.

[pierrejeambrun]

The only issue I see with that is:

A query will be emitted each time we do:

dag_run.dag_versions()

When we will serialize multiple dag_runs (100 of them), this will happen in a loop, resulting in an extra 100 query being executed to fetch dag_versions for each dag_run.

What would be great is the ability to query dag_run and eagerly load that: (pseudo code)

query(DagRun).options(joinedload(DagRun.dag_versions))

This way in the webserver we can eagerly load all the dag_versions of the dag_run query all at the same time in one query.

[ephraimbuddy]

joinedload won't be possible since we want only the IDs from DR.tis

[jedcunningham]

Might be able to do something like this:

query(DagRun).options(
joinedload(DagRun.tis)
joinedload(TaskInstance.dag_versions)
joinedload(DagRun.tihs)
joinedload(TaskInstanceHistory.dag_versions)
)

Basically, joinedload both ti and this, and their respective dag_versions. Not sure if that'd work, it might.

[pierrejeambrun]

I’m not sure if this will have any effect on the “dag_versions” method. Counting on caching to not emit the following requests ?

Maybe we can transform that into a relationship or a hybrid property so we can use loading options / explicitly select them ?

[jedcunningham]

Looking more closely, if we just add a tih relationship to dagrun, we should be able to joinedload them - and we don't need the whole version record, just the ids are enough, and those are already on ti/tih.

But, instead of having a method on dagrun to do this, it might be better to have a helper/classmethod that knows how to do this in bulk, e.g. def get_dag_version_ids(list[DagRun], dict[DagRun, list[uuids]]).

[jason810496]

How about using cached_property instead of the property decorator and using union without joinload support? Would this make it much simpler?

    @cache_property
    @provide_session
    def dag_versions(self, session: Session = NEW_SESSION) -> list[DagVersion]:
        """
        Return the DAG versions associated with the TIs of this DagRun.
        :param session: database session
        """
        tis = select(TI.dag_version_id).where(
            TI.dag_id == self.run_id, TI.run_id == self.run_id
        )
        tih = select(TIH.dag_version_id).where(
            TIH.dag_id == self.run_id,
            TIH.run_id == self.run_id,
        )
        return (
            session.execute(
                tis.distinct().union(tih.distinct())
            )
            .scalars()
            .all()
        )

[ephraimbuddy]

@jedcunningham That will still make a separate query when accessed. I will suggest we keep what we have currently and add the above in a separate PR.

[ephraimbuddy]

@pierrejeambrun , won't it be easier to retain the DagVersion.dag_version_id and filter by the TI.dag_version_id in API response?

[ephraimbuddy]

@jason810496 , your query is good, though it will still make a separate query when accessed, and I'm not sure how cached_property would help since we might not access it repeatedly

[pierrejeambrun]

Nice!

[pierrejeambrun]

The only issue I see with that is:

A query will be emitted each time we do:

dag_run.dag_versions()

When we will serialize multiple dag_runs (100 of them), this will happen in a loop, resulting in an extra 100 query being executed to fetch dag_versions for each dag_run.

What would be great is the ability to query dag_run and eagerly load that: (pseudo code)

query(DagRun).options(joinedload(DagRun.dag_versions))

This way in the webserver we can eagerly load all the dag_versions of the dag_run query all at the same time in one query.

[pierrejeambrun]

Is that related to this PR ?

[ephraimbuddy]

No. Had issues with pre-commit. Will remove it when ready