rjgoyln commented on code in PR #72185:
URL: https://github.com/apache/airflow/pull/72185#discussion_r4187024501
##########
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:
One small test gap: this test also passes on `main`, since the per-TI
`set_state` path produces the same states, dates, and durations.
Would it make sense to use `assert_queries_count(5)` here to pin the single
flush and catch a regression to per-TI flushing?
```suggestion
with assert_queries_count(5):
updated_tis, killed_tis = set_dag_run_state_to_failed(
dag=dag, run_id=dr.run_id, commit=True, session=session
)
```
##########
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:
Could we also pin the state of finished siblings here? I noticed that on
main, they can be overwritten as well: a finished success sibling becomes
failed when marking the run failed, and a failed sibling becomes success when
marking it successful. With this change, they keep their existing state as
expected.
Would it make sense to add a fourth map index that is already success and
assert that it remains success?
--
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]
