Lee-W commented on code in PR #71072:
URL: https://github.com/apache/airflow/pull/71072#discussion_r4216932329
##########
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):
Review Comment:
Now using `_make_n_consumer_dags_each_with_one_pending_apdr`. Moved
`_count_audit_log` and `_get_cap_reached_levels` as module-level shared helpers.
--
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]
