dgvj-work commented on code in PR #71459:
URL: https://github.com/apache/airflow/pull/71459#discussion_r4509954892
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py:
##########
@@ -144,9 +143,7 @@
)
def get_dag_run(dag_id: str, dag_run_id: str, session: SessionDep) ->
DAGRunResponse:
dag_run = session.scalar(
- select(DagRun)
- .filter_by(dag_id=dag_id, run_id=dag_run_id)
- .options(joinedload(DagRun.dag_model),
*eager_load_teams(DagRun.dag_model))
+ select(DagRun).filter_by(dag_id=dag_id,
run_id=dag_run_id).options(*eager_load_dag_run_for_list())
Review Comment:
Reverted. `get_dag_run()` is back to `joinedload(DagRun.dag_model)` plus
`eager_load_teams`, so it no longer joins `created_dag_version` and its bundle.
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py:
##########
@@ -143,6 +160,76 @@ def custom_timetable_plugin(monkeypatch):
DAG_RUNS_LIST = [DAG1_RUN1_ID, DAG1_RUN2_ID, DAG2_RUN1_ID, DAG2_RUN2_ID]
+def _create_bundled_run_with_history(dag_id, dag_maker, session):
+ with dag_maker(
+ dag_id=dag_id, schedule=None, start_date=START_DATE1,
bundle_version="v1", serialized=True
+ ):
+ EmptyOperator(task_id="0")
+ EmptyOperator(task_id="1")
+ dag_run = dag_maker.create_dagrun(state=DagRunState.SUCCESS)
+ old_dag_version = DagVersion.get_latest_version(dag_id)
+ for ti in dag_run.task_instances:
+ ti.state = TaskInstanceState.SUCCESS
+ session.merge(ti)
+ session.flush()
+
+ with dag_maker(
+ dag_id=dag_id, schedule=None, start_date=START_DATE1,
bundle_version="v2", serialized=True
+ ):
+ EmptyOperator(task_id="0")
+ EmptyOperator(task_id="1")
+ new_dag_version = DagVersion.get_latest_version(dag_id)
+ assert old_dag_version.id != new_dag_version.id
+
+ clear_task_instances(list(dag_run.task_instances), session,
run_on_latest_version=True)
+ session.commit()
+
+ archived_version_ids = set(
+ session.scalars(
+ select(TaskInstance.dag_version_id)
+ .where(
+ TaskInstance.dag_id == dag_id,
+ TaskInstance.run_id == dag_run.run_id,
+ TaskInstance.dag_version_id.is_not(None),
+ )
+ .execution_options(include_all_attempts=True)
+ )
+ )
+ assert old_dag_version.id in archived_version_ids
+ return dag_run.run_id, old_dag_version, new_dag_version
+
+
+def _dag_bag_from_serialized(session):
+ """DagBag that resolves DAGs from SerializedDagModel in the current
session."""
+ dag_bag = mock.Mock()
+
+ def get_dag_for_run(dag_run, session=session):
+ sdm = None
+ if dag_run.created_dag_version_id is not None:
+ sdm = session.scalar(
+ select(SerializedDagModel).where(
+ SerializedDagModel.dag_version_id ==
dag_run.created_dag_version_id
+ )
+ )
+ if sdm is None:
+ sdm = SerializedDagModel.get(dag_run.dag_id, session=session)
+ return None if sdm is None else sdm.dag
+
+ dag_bag.get_dag_for_run.side_effect = get_dag_for_run
+ return dag_bag
+
+
+def _assert_serialization_leaves_ti_collections_unloaded(dag_run, *,
old_dag_version, new_dag_version):
+ session = sa_inspect(dag_run).session
+ session.expire(dag_run, ["task_instances", "historical_task_instances"])
Review Comment:
Dropped `_assert_serialization_leaves_ti_collections_unloaded`, and the
PATCH, single-clear, and trigger tests with it. GET and bulk clear still assert
the response includes both the current version and the archived-attempt version.
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -1750,16 +1752,169 @@ def
test_dag_run_dag_versions_with_null_created_dag_version(self, dag_maker, ses
EmptyOperator(task_id="empty")
dag_run = dag_maker.create_dagrun()
+ ti_version_ids = {ti.dag_version_id for ti in dag_run.task_instances
if ti.dag_version_id is not None}
+ assert ti_version_ids
+
dag_run.bundle_version = "some_bundle_version"
dag_run.created_dag_version_id = None
dag_run.created_dag_version = None
session.merge(dag_run)
session.flush()
- # This should return empty list, not [None]
- assert dag_run.dag_versions == []
- assert isinstance(dag_run.dag_versions, list)
- assert len(dag_run.dag_versions) == 0
+ # Derive from TI versions; never return [None] from a null
created_dag_version shortcut.
+ versions = dag_run.dag_versions
+ assert isinstance(versions, list)
+ assert {dv.id for dv in versions} == ti_version_ids
+ assert None not in versions
+
+ def
test_bundled_dag_run_dag_versions_after_partial_run_on_latest_version(self,
dag_maker, session):
+ """Bundled dag_versions include uncleared TI versions after a partial
run_on_latest_version clear.
+
+ Clearing a subset with run_on_latest_version=True bumps
created_dag_version_id
+ and bundle_version to latest while uncleared TIs stay on the old
version.
+ The property must report both.
+ """
+ with dag_maker(
+ "test_dag_run_dag_versions_partial_latest",
+ schedule=datetime.timedelta(days=1),
+ start_date=DEFAULT_DATE,
+ bundle_version="v1",
+ ):
+ EmptyOperator(task_id="0")
+ EmptyOperator(task_id="1")
+ dag_run = dag_maker.create_dagrun(state=State.RUNNING,
run_type=DagRunType.SCHEDULED)
+
+ old_dag_version = DagVersion.get_latest_version(dag_run.dag_id)
+ ti0, ti1 = sorted(dag_run.task_instances, key=lambda ti: ti.task_id)
+ for ti in (ti0, ti1):
+ ti.state = TaskInstanceState.SUCCESS
+ session.merge(ti)
+ dag_run.state = DagRunState.SUCCESS
+ session.merge(dag_run)
+ session.flush()
+
+ with dag_maker(
+ "test_dag_run_dag_versions_partial_latest",
+ schedule=datetime.timedelta(days=1),
+ start_date=DEFAULT_DATE,
+ bundle_version="v2",
+ ):
+ EmptyOperator(task_id="0")
+ EmptyOperator(task_id="1")
+ new_dag_version = DagVersion.get_latest_version(dag_run.dag_id)
+ assert old_dag_version.id != new_dag_version.id
+
+ clear_task_instances([ti0], session, run_on_latest_version=True)
+ session.commit()
+
+ dag_run = session.get(DagRun, dag_run.id)
+ assert dag_run is not None
+ assert dag_run.created_dag_version_id == new_dag_version.id
+ assert dag_run.bundle_version == new_dag_version.bundle_version
+
+ tis = {ti.task_id: ti for ti in dag_run.task_instances}
+ assert tis["0"].dag_version_id == new_dag_version.id
+ assert tis["1"].dag_version_id == old_dag_version.id
+
+ session.expire(dag_run, ["task_instances",
"historical_task_instances"])
+ dag_run._prefetched_dag_version_ids = None
+ unloaded = sa_inspect(dag_run).unloaded
+ assert "task_instances" in unloaded
+ assert "historical_task_instances" in unloaded
+ version_ids = {dv.id for dv in dag_run.dag_versions}
+ assert version_ids == {old_dag_version.id, new_dag_version.id}
+ assert "task_instances" in sa_inspect(dag_run).unloaded
+ assert "historical_task_instances" in sa_inspect(dag_run).unloaded
+
+ attach_dag_versions_to_runs([dag_run], session=session)
+ assert dag_run._prefetched_dag_version_ids is not None
+ assert set(dag_run._prefetched_dag_version_ids) ==
{old_dag_version.id, new_dag_version.id}
+ assert {dv.id for dv in dag_run.dag_versions} == {old_dag_version.id,
new_dag_version.id}
+
+ def
test_bundled_dag_run_dag_versions_after_full_run_on_latest_version(self,
dag_maker, session):
+ """Bundled dag_versions still include archived-attempt versions after
a full clear.
+
+ Clearing every TI with run_on_latest_version=True bumps
created_dag_version_id
+ and bundle_version to latest, and all current TIs move to that
version. The
+ previous version remains on archived attempts, so the property must
still report both.
+ """
+ with dag_maker(
+ "test_dag_run_dag_versions_full_latest",
+ schedule=datetime.timedelta(days=1),
+ start_date=DEFAULT_DATE,
+ bundle_version="v1",
+ ):
+ EmptyOperator(task_id="0")
+ EmptyOperator(task_id="1")
+ dag_run = dag_maker.create_dagrun(state=State.RUNNING,
run_type=DagRunType.SCHEDULED)
+
+ old_dag_version = DagVersion.get_latest_version(dag_run.dag_id)
+ for ti in dag_run.task_instances:
+ ti.state = TaskInstanceState.SUCCESS
+ session.merge(ti)
+ dag_run.state = DagRunState.SUCCESS
+ session.merge(dag_run)
+ session.flush()
+
+ with dag_maker(
+ "test_dag_run_dag_versions_full_latest",
+ schedule=datetime.timedelta(days=1),
+ start_date=DEFAULT_DATE,
+ bundle_version="v2",
+ ):
+ EmptyOperator(task_id="0")
+ EmptyOperator(task_id="1")
+ new_dag_version = DagVersion.get_latest_version(dag_run.dag_id)
+ assert old_dag_version.id != new_dag_version.id
+
+ clear_task_instances(list(dag_run.task_instances), session,
run_on_latest_version=True)
+ session.commit()
+
+ dag_run = session.get(DagRun, dag_run.id)
+ assert dag_run is not None
+ assert dag_run.created_dag_version_id == new_dag_version.id
+ assert dag_run.bundle_version == new_dag_version.bundle_version
+ assert {ti.dag_version_id for ti in dag_run.task_instances} ==
{new_dag_version.id}
+
+ archived_version_ids = set(
+ session.scalars(
+ select(TaskInstance.dag_version_id)
+ .where(
+ TaskInstance.dag_id == dag_run.dag_id,
+ TaskInstance.run_id == dag_run.run_id,
+ TaskInstance.dag_version_id.is_not(None),
+ )
+ .execution_options(include_all_attempts=True)
+ )
+ )
+ assert old_dag_version.id in archived_version_ids
+
+ session.expire(dag_run, ["task_instances",
"historical_task_instances"])
+ dag_run._prefetched_dag_version_ids = None
+ unloaded = sa_inspect(dag_run).unloaded
+ assert "task_instances" in unloaded
+ assert "historical_task_instances" in unloaded
+ version_ids = {dv.id for dv in dag_run.dag_versions}
+ assert version_ids == {old_dag_version.id, new_dag_version.id}
+ assert "task_instances" in sa_inspect(dag_run).unloaded
+ assert "historical_task_instances" in sa_inspect(dag_run).unloaded
+
+ attach_dag_versions_to_runs([dag_run], session=session)
+ assert dag_run._prefetched_dag_version_ids is not None
+ assert set(dag_run._prefetched_dag_version_ids) ==
{old_dag_version.id, new_dag_version.id}
+ assert {dv.id for dv in dag_run.dag_versions} == {old_dag_version.id,
new_dag_version.id}
+
+ def test_dag_run_dag_versions_transient_run_uses_relationships(self):
Review Comment:
Set `bundle_version` on the transient run. The old `created_dag_version`
shortcut would return an empty list here, so the test fails if that property
change is reverted.
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -598,22 +599,55 @@ def dag_run_data(self) -> DRDataModel:
partition_date=self.partition_date,
)
+ @classmethod
+ def fetch_dag_versions_for_runs(
+ cls, run_keys: Collection[tuple[str, str]], *, session: Session
+ ) -> dict[tuple[str, str], dict[UUID, DagVersion]]:
+ """Distinct DagVersions per (dag_id, run_id) across current and
archived attempts, one query."""
Review Comment:
The distinct ids and the DagVersion load, including bundle, are now one
statement: a DISTINCT subquery joined to DagVersion. Distinct stays on the id
because `version_data` is JSON. The `dag_versions` docstring now says Dag.
##########
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py:
##########
@@ -1709,6 +1818,34 @@ def test_should_respond_403(self,
unauthorized_test_client):
response = unauthorized_test_client.patch("/dags/dag_1/dagRuns/run_1",
json={})
assert response.status_code == 403
+ @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle")
+ @mock.patch(
+ "airflow.api_fastapi.core_api.routes.public.dag_run.get_dag_for_run",
+ autospec=True,
+ return_value=mock.Mock(),
Review Comment:
Removed the PATCH and trigger mocks with those tests.
`_dag_bag_from_serialized` now uses `mock.Mock(spec=DBDagBag)`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
