Lee-W commented on code in PR #71072:
URL: https://github.com/apache/airflow/pull/71072#discussion_r4256554618
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -14345,6 +14415,617 @@ 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": PARTITION_CAP_REACHED_EVENT,
+ "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 PARTITION_CAP_REACHED_EVENT 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 PARTITION_CAP_REACHED_EVENT 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")
[email protected]("attr", ["is_stale", "is_paused", "is_draining"])
+def test_partition_cap_reporting_excludes_inactive_dag_decoy(
+ attr: str, 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 Dag-state predicates: ``is_stale``, ``is_paused`` and
``is_draining`` must all be
+ ``False``. This test adds a decoy pending APDR whose target Dag became
inactive after the
+ APDR was written -- a count query missing that filter would wrongly
include it,
+ over-counting the backlog even though the fetch query never selects it.
+ """
+ cap = 3
+ suffix = attr.replace("_", "-")
+ _make_n_satisfied_apdrs(
+ consumer_dag_id=f"cap-consumer-inactive-decoy-{suffix}",
+ asset=Asset(name=f"asset-cap-inactive-decoy-{suffix}"),
+ partition_keys=["k1", "k2", "k3"],
+ session=session,
+ dag_maker=dag_maker,
+ )
+
+ inactive_consumer_dag_id = f"cap-consumer-inactive-decoy-{suffix}-inactive"
+ inactive_asset = Asset(name=f"asset-cap-inactive-decoy-{suffix}-inactive")
+ with dag_maker(
+ dag_id=inactive_consumer_dag_id,
+ schedule=PartitionedAssetTimetable(assets=inactive_asset,
default_partition_mapper=IdentityMapper()),
+ session=session,
+ ):
+ EmptyOperator(task_id="hi")
+ session.commit()
+
+ _produce_and_register_asset_event(
+ dag_id=f"asset-event-producer-inactive-decoy-{suffix}",
+ asset=inactive_asset,
+ partition_key="decoy",
+ session=session,
+ dag_maker=dag_maker,
+ )
+
+ # Flip the flag only after the APDR exists: AssetManager drops paused and
draining Dags
+ # before writing APDRs, so flag-first would produce no decoy at all.
+ dm = session.get(DagModel, inactive_consumer_dag_id)
+ assert dm is not None
+ setattr(dm, attr, 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 PARTITION_CAP_REACHED_EVENT 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.
+ """
+ apdrs = _make_n_consumer_dags_each_with_one_pending_apdr(
+ name_prefix="many", n=7, 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:02d}" for i in range(1,
8))
+ assert {
+ "event": PARTITION_CAP_REACHED_EVENT,
+ "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
Review Comment:
Removed. Thanks!
--
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]
