matthieuauger commented on code in PR #72185:
URL: https://github.com/apache/airflow/pull/72185#discussion_r4258786673
##########
airflow-core/tests/unit/api/common/test_mark_tasks.py:
##########
@@ -61,6 +64,75 @@ def test_set_dag_run_state_to_failed(dag_maker:
DagMaker[SerializedDAG]):
assert "teardown" not in task_dict
+def test_set_dag_run_state_to_failed_skips_pending_tis_in_bulk(dag_maker:
DagMaker[SerializedDAG]):
+ """Pending TIs that never started are skipped by one UPDATE; one that did
keeps its start_date."""
+ with dag_maker("TEST_DAG_1") as dag:
+ EmptyOperator(task_id="never_started")
+ EmptyOperator(task_id="never_started_2")
+ EmptyOperator(task_id="up_for_retry")
+ dr = dag_maker.create_dagrun()
+ session = dag_maker.session
+ started_at = timezone.datetime(2024, 1, 1, 12, 0, 0)
+ for ti in dr.get_task_instances(session=session):
+ if ti.task_id == "up_for_retry":
+ ti.state = TaskInstanceState.UP_FOR_RETRY
+ ti.start_date = started_at
+ session.flush()
+
+ updated_tis, killed_tis = set_dag_run_state_to_failed(
+ dag=dag, run_id=dr.run_id, commit=True, session=session
+ )
Review Comment:
Good point, thanks — the test did pass on main. Done in a0962a7898: the call
now runs under `assert_queries_count(5)` — the two SELECTs, the SELECT of the
dag run, and one UPDATE per set of changed columns (the never-started TIs also
get a `start_date`, so two batches for the three TIs). On main the same test
counts 6, one UPDATE per TI, so the cap catches a regression to per-TI
flushing. Also rebased on main.
##########
airflow-core/tests/unit/api/common/test_mark_tasks.py:
##########
@@ -61,6 +64,75 @@ def test_set_dag_run_state_to_failed(dag_maker:
DagMaker[SerializedDAG]):
assert "teardown" not in task_dict
+def test_set_dag_run_state_to_failed_skips_pending_tis_in_bulk(dag_maker:
DagMaker[SerializedDAG]):
+ """Pending TIs that never started are skipped by one UPDATE; one that did
keeps its start_date."""
+ with dag_maker("TEST_DAG_1") as dag:
+ EmptyOperator(task_id="never_started")
+ EmptyOperator(task_id="never_started_2")
+ EmptyOperator(task_id="up_for_retry")
+ dr = dag_maker.create_dagrun()
+ session = dag_maker.session
+ started_at = timezone.datetime(2024, 1, 1, 12, 0, 0)
+ for ti in dr.get_task_instances(session=session):
+ if ti.task_id == "up_for_retry":
+ ti.state = TaskInstanceState.UP_FOR_RETRY
+ ti.start_date = started_at
+ session.flush()
+
+ updated_tis, killed_tis = set_dag_run_state_to_failed(
+ dag=dag, run_id=dr.run_id, commit=True, session=session
+ )
+ session.flush()
+
+ assert killed_tis == []
+ assert {ti.task_id for ti in updated_tis} == {"never_started",
"never_started_2", "up_for_retry"}
+ # The returned objects reflect the new state without a reload...
+ assert all(ti.state == TaskInstanceState.SKIPPED for ti in updated_tis)
+ # ...and so does the database.
+ session.expire_all()
+ task_dict = {ti.task_id: ti for ti in
dr.get_task_instances(session=session)}
+ for task_id in ("never_started", "never_started_2"):
+ ti = task_dict[task_id]
+ assert ti.state == TaskInstanceState.SKIPPED
+ assert ti.start_date is not None
+ assert ti.end_date == ti.start_date
+ assert ti.duration == 0
+ ti = task_dict["up_for_retry"]
+ assert ti.state == TaskInstanceState.SKIPPED
+ assert ti.start_date == started_at
+ assert ti.end_date is not None
+ assert ti.duration == pytest.approx((ti.end_date -
started_at).total_seconds())
+
+
+def
test_set_dag_run_state_to_failed_mapped_task_only_fails_running_map_indexes(
+ dag_maker: DagMaker[SerializedDAG],
+):
+ """Only the running map indexes fail; the pending siblings of the same
mapped task are skipped."""
+ with dag_maker("TEST_DAG_1") as dag:
+ MockOperator.partial(task_id="mapped").expand(arg2=[1, 2, 3])
Review Comment:
Done in a0962a7898: the task now expands to four map indexes, with index 3
already SUCCESS before the call; it is neither in `updated_tis` nor
overwritten. I checked the other direction too: against main's code, the
bare-task `set_state` selected every map index and index 3 ended up FAILED.
--
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]
