Lee-W commented on code in PR #71072:
URL: https://github.com/apache/airflow/pull/71072#discussion_r4216935419
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -13244,6 +13296,594 @@ def
test_partition_cap_at_n_minus_one_leaves_one_pending(dag_maker: DagMaker, se
assert partition_dags == {"cap-consumer-n-minus-one"}
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
[email protected](
+ ("cap", "partition_keys", "expect_cap_hit"),
+ [
+ pytest.param(2, ["k1", "k2", "k3"], True, id="over-cap"),
+ pytest.param(3, ["k1", "k2"], False, id="under-cap"),
+ # Regression guard for the ``backlog_total > cap`` boundary:
``backlog_total == cap``
+ # must not be misreported as a backlog (that would be the "exactly
cap, nothing left"
+ # case wrongly reported as "more than cap, backlog remains").
+ pytest.param(3, ["k1", "k2", "k3"], False, id="exactly-cap"),
+ ],
+)
+def test_partition_cap_reporting(
+ dag_maker: DagMaker,
+ session: Session,
+ caplog,
+ cap: int,
+ partition_keys: list[str],
+ expect_cap_hit: bool,
+):
+ """Only a pending count strictly above the cap reports a backlog, via log
and audit `Log` row."""
+ suffix = f"{cap}-{len(partition_keys)}"
+ _make_n_satisfied_apdrs(
+ consumer_dag_id=f"cap-consumer-{suffix}",
+ asset=Asset(name=f"asset-cap-{suffix}"),
+ partition_keys=partition_keys,
+ session=session,
+ dag_maker=dag_maker,
+ )
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ runner._max_partition_dag_runs_per_loop = cap
+ runner._create_dagruns_for_partitioned_asset_dags(session=session)
+
+ dag_id = f"cap-consumer-{suffix}"
+ audit_rows = session.scalars(select(Log).where(Log.event == "partition Dag
run cap reached")).all()
+ if expect_cap_hit:
+ assert {
+ "event": "Reached the per-tick cap on pending partitioned Dag
runs; the remaining backlog "
+ "will be evaluated over subsequent scheduler ticks",
+ "cap": cap,
+ "backlog_total": 3,
+ "dag_ids": [dag_id],
+ } in caplog
+ assert [row.event for row in audit_rows] == ["partition Dag run cap
reached"]
+ extra = audit_rows[0].extra
+ assert extra is not None
+ assert f"Affected dag_ids (highest pending count first): {dag_id}." in
extra
+ else:
+ assert "Reached the per-tick cap on pending partitioned Dag runs" not
in caplog
+ assert audit_rows == []
+
+
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
+def test_partition_cap_reporting_excludes_already_fired_decoy(dag_maker:
DagMaker, session: Session, caplog):
+ """
+ The ``backlog_total`` count query's ``WHERE`` clause must mirror the main
query's.
+
+ ``exactly-cap`` in :func:`test_partition_cap_reporting` cannot catch a
filter drift in the
+ count query (e.g. a dropped ``created_dag_run_id.is_(None)`` filter):
every row in that
+ table at that point is a genuine, unfired, non-stale APDR, so a looser
filter has nothing
+ extra to over-count. This test adds a decoy APDR that already fired
+ (``created_dag_run_id`` set) — a naive/unfiltered ``count()`` would
wrongly include it.
+ """
+ cap = 3
+ asset = Asset(name="asset-cap-decoy")
+ _make_n_satisfied_apdrs(
+ consumer_dag_id="cap-consumer-decoy",
+ asset=asset,
+ partition_keys=["k1", "k2", "k3"],
+ session=session,
+ dag_maker=dag_maker,
+ )
+ decoy = _produce_and_register_asset_event(
+ dag_id="asset-event-producer-decoy",
+ asset=asset,
+ partition_key="decoy",
+ session=session,
+ dag_maker=dag_maker,
+ )
+ fired_dag_run_id =
session.scalar(select(DagRun.id).order_by(DagRun.id.desc()).limit(1))
+ assert fired_dag_run_id is not None
+ decoy.created_dag_run_id = fired_dag_run_id
+ session.commit()
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ runner._max_partition_dag_runs_per_loop = cap
+ runner._create_dagruns_for_partitioned_asset_dags(session=session)
+
+ assert "Reached the per-tick cap on pending partitioned Dag runs" not in
caplog
+ audit_rows = session.scalars(select(Log).where(Log.event == "partition Dag
run cap reached")).all()
+ assert audit_rows == []
+
+
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
+def test_partition_cap_reporting_excludes_stale_dag_decoy(dag_maker: DagMaker,
session: Session, caplog):
+ """
+ The ``backlog_total`` count query's ``WHERE`` clause must mirror the main
query's.
+
+ Mirrors :func:`test_partition_cap_reporting_excludes_already_fired_decoy`
for the fetch
+ query's other predicate: ``DagModel.is_stale.is_(False)``. This test adds
a decoy pending
+ APDR whose target Dag has since gone stale (removed from its Dag file) — a
count query
+ missing that filter would wrongly include it, over-counting the backlog
even though the
+ fetch query (unaffected by this specific drift) never selects it.
+ """
+ cap = 3
+ _make_n_satisfied_apdrs(
+ consumer_dag_id="cap-consumer-stale-decoy",
+ asset=Asset(name="asset-cap-stale-decoy"),
+ partition_keys=["k1", "k2", "k3"],
+ session=session,
+ dag_maker=dag_maker,
+ )
+
+ stale_consumer_dag_id = "cap-consumer-stale-decoy-stale"
+ stale_asset = Asset(name="asset-cap-stale-decoy-stale")
+ with dag_maker(
+ dag_id=stale_consumer_dag_id,
+ schedule=PartitionedAssetTimetable(assets=stale_asset,
default_partition_mapper=IdentityMapper()),
+ session=session,
+ ):
+ EmptyOperator(task_id="hi")
+ session.commit()
+
+ _produce_and_register_asset_event(
+ dag_id="asset-event-producer-stale-decoy",
+ asset=stale_asset,
+ partition_key="decoy",
+ session=session,
+ dag_maker=dag_maker,
+ )
+
+ dm = session.get(DagModel, stale_consumer_dag_id)
+ assert dm is not None
+ dm.is_stale = True
+ session.commit()
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ runner._max_partition_dag_runs_per_loop = cap
+ runner._create_dagruns_for_partitioned_asset_dags(session=session)
+
+ assert "Reached the per-tick cap on pending partitioned Dag runs" not in
caplog
+ audit_rows = session.scalars(select(Log).where(Log.event == "partition Dag
run cap reached")).all()
+ assert audit_rows == []
+
+
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
+def test_partition_cap_reporting_lists_all_affected_dag_ids(dag_maker:
DagMaker, session: Session, caplog):
+ """
+ The log/audit ``dag_ids`` list surfaces every distinct dag_id in the
backlog, not just this
+ tick's oldest-cap slice.
+
+ Seven consumer Dags each get one satisfied APDR; with the cap at 6,
`pending_apdrs` (this
+ tick's oldest six) spans only six distinct dag_ids, but the seventh Dag's
APDR is still part
+ of the backlog `count()` sees — its dag_id must appear in both the log
event and the audit
+ row too, even though its Dag run won't be created until the next tick.
+
+ Built inline rather than via :func:`_make_n_satisfied_apdrs` because that
helper's producer
+ dag_id numbering restarts at 1 on every call, colliding across these seven
consumer Dags.
+ """
+ apdrs = []
+ for i in range(1, 8):
+ asset = Asset(name=f"asset-cap-many-{i}")
+ with dag_maker(
+ dag_id=f"cap-consumer-many-{i}",
+ schedule=PartitionedAssetTimetable(assets=asset,
default_partition_mapper=IdentityMapper()),
+ session=session,
+ ):
+ EmptyOperator(task_id="hi")
+ session.commit()
+ apdrs.append(
+ _produce_and_register_asset_event(
+ dag_id=f"asset-event-producer-many-{i}",
+ asset=asset,
+ partition_key=f"k{i}",
+ session=session,
+ dag_maker=dag_maker,
+ )
+ )
+ base = timezone.utcnow()
+ for i, apdr in enumerate(apdrs):
+ apdr.created_at = base + timedelta(seconds=i)
+ session.commit()
+
+ runner = SchedulerJobRunner(
+ job=Job(job_type=SchedulerJobRunner.job_type),
executors=[MockExecutor(do_update=False)]
+ )
+ runner._max_partition_dag_runs_per_loop = 6
+ runner._create_dagruns_for_partitioned_asset_dags(session=session)
+
+ # The dag_ids list is sorted by (-count, dag_id); all seven Dags tie on
count (one pending
+ # APDR each), so the tie-break is dag_id ascending — a deterministic,
reproducible order,
+ # not just membership. All seven Dags are expected, including the seventh
whose APDR is
+ # deferred past this tick's cap.
+ expected_dag_ids = sorted(f"cap-consumer-many-{i}" for i in range(1, 8))
+ assert {
+ "event": "Reached the per-tick cap on pending partitioned Dag runs;
the remaining backlog "
+ "will be evaluated over subsequent scheduler ticks",
+ "cap": 6,
+ "backlog_total": 7,
+ "dag_ids": expected_dag_ids,
+ } in caplog
+
+ audit_row = session.scalar(select(Log).where(Log.event == "partition Dag
run cap reached"))
+ assert audit_row is not None
+ assert audit_row.extra is not None
+ note_match = re.search(r"Affected dag_ids \(highest pending count first\):
(.+)\.$", audit_row.extra)
+ assert note_match is not None
+ assert note_match.group(1).split(", ") == expected_dag_ids
+ assert "and 1 more" not in audit_row.extra
+
+
[email protected]_serialized_dag
[email protected]("clear_asset_partition_rows", "clear_audit_log_rows")
+def test_partition_cap_reporting_truncates_dag_id_list_beyond_max(
+ dag_maker: DagMaker, session: Session, caplog
+):
+ """
+ The log/audit ``dag_ids`` list is capped at
``MAX_PARTITION_CAP_BACKLOG_DAG_IDS_LOGGED``,
+ ordered by descending per-dag pending count with a dag_id tie-break — not
alphabetically
+ truncated — and the audit row's prose notes how many dag_ids were dropped
by the cap.
+
+ Built via :func:`_make_n_consumer_dags_each_with_one_pending_apdr` rather
than
+ :func:`_make_n_satisfied_apdrs` because that helper's producer dag_id
numbering restarts at 1
+ on every call, colliding across these 21 consumer Dags — same issue
documented on
+ :func:`test_partition_cap_reporting_lists_all_affected_dag_ids`.
+ """
+ _make_n_consumer_dags_each_with_one_pending_apdr(
+ name_prefix="trunc-low", n=20, session=session, dag_maker=dag_maker
+ )
+
+ high_asset = Asset(name="asset-cap-trunc-high")
+ with dag_maker(
+ dag_id="cap-consumer-trunc-high",
Review Comment:
Renamed to `cap-consumer-trunc-zz-high`. It now sorts after the `low` Dags
by name, so the test only passes with the descending-count order.
--
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]
