ephraimbuddy commented on code in PR #71459:
URL: https://github.com/apache/airflow/pull/71459#discussion_r4455768449
##########
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:
These made sense when each route had to opt in to the prefetch. Now this
helper clears the prefetch and expires both collections before serializing, so
every route test only exercises `dag_versions` itself, which the model tests
already cover. I'd keep the response-level assertions in the GET and bulk clear
tests, and drop the PATCH, single clear and trigger tests along with this
helper.
##########
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:
nit: These mocks need a `spec`, and so do the new ones in
`_dag_bag_from_serialized` and the trigger test's `request`.
##########
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:
nit: This isn't needed for the fix anymore, and it joins
`created_dag_version` and its bundle, which `DAGRunResponse` doesn't serialize.
I'd revert it.
##########
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:
nit: This passes with the property change reverted, because a transient run
without `bundle_version` already used the relationships on main. Set
`bundle_version` on the run so it covers the change, or drop it.
##########
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:
nit: This is two queries (the DISTINCT plus the DagVersion load). Also, the
`dag_versions` docstring below should say Dag, not DAG.
--
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]
