Skip to content

Fix dataset-triggered DAGs firing without all upstream events #63939

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
joseotaviorf wants to merge 1 commit into from
Closed

Fix dataset-triggered DAGs firing without all upstream events #63939

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions airflow/datasets/manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
DatasetModel,
)
from airflow.stats import Stats
from airflow.utils import timezone
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import NEW_SESSION, provide_session

Expand Down Expand Up @@ -184,7 +185,7 @@ def _slow_path_queue_dagruns(
cls, dataset_id: int, dags_to_queue: set[DagModel], session: Session
) -> None:
def _queue_dagrun_if_needed(dag: DagModel) -> str | None:

@kaxil kaxil Mar 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Postgres path has a WHERE guard (created_at < excluded.created_at) that prevents a slower transaction from overwriting a newer timestamp, but session.merge() here still overwrites unconditionally. Two concurrent tasks can race: if the slower one commits last, it writes an older created_at. Using utcnow() makes backward drift unlikely, but it's still possible under clock adjustments or NTP drift.

item = DatasetDagRunQueue(target_dag_id=dag.dag_id, dataset_id=dataset_id)
item = DatasetDagRunQueue(target_dag_id=dag.dag_id, dataset_id=dataset_id, created_at=timezone.utcnow())
# Don't error whole transaction when a single RunQueue item conflicts.
# https://docs.sqlalchemy.org/en/14/orm/session_transaction.html#using-savepoint
try:
Expand All @@ -202,8 +203,13 @@ def _queue_dagrun_if_needed(dag: DagModel) -> str | None:
def _postgres_queue_dagruns(cls, dataset_id: int, dags_to_queue: set[DagModel], session: Session) -> None:
from sqlalchemy.dialects.postgresql import insert

values = [{"target_dag_id": dag.dag_id} for dag in dags_to_queue]
stmt = insert(DatasetDagRunQueue).values(dataset_id=dataset_id).on_conflict_do_nothing()
values = [{"target_dag_id": dag.dag_id, "created_at": timezone.utcnow()} for dag in dags_to_queue]

@kaxil kaxil Mar 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Each iteration calls timezone.utcnow() separately, so each DAG in dags_to_queue gets a slightly different timestamp. For the common single-DAG case this doesn't matter, but with multiple DAGs the timestamps diverge. Computing the timestamp once before the comprehension would be more consistent:

now = timezone.utcnow()
values = [{"target_dag_id": dag.dag_id, "created_at": now} for dag in dags_to_queue]

stmt = insert(DatasetDagRunQueue).values(dataset_id=dataset_id)
stmt = stmt.on_conflict_do_update(
index_elements=["dataset_id", "target_dag_id"],
set_={"created_at": stmt.excluded.created_at},
where=(DatasetDagRunQueue.created_at < stmt.excluded.created_at),
)
session.execute(stmt, values)
Comment on lines +206 to 213

Copilot AI Apr 2, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In _postgres_queue_dagruns, values = [{..., "created_at": timezone.utcnow()} for dag in dags_to_queue] evaluates utcnow() once per DAG. That means created_at can differ slightly across rows from the same dataset event, and it adds unnecessary per-row Python work. Consider computing a single now = timezone.utcnow() once and reusing it for all rows in this batch to keep timestamps consistent and reduce overhead.

Copilot uses AI. Check for mistakes.

@classmethod
Expand Down
32 changes: 18 additions & 14 deletions airflow/jobs/scheduler_job_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1489,21 +1489,25 @@ def _create_dag_runs_dataset_triggered(
events=dataset_events,
)

dag_run = dag.create_dagrun(
run_id=run_id,
run_type=DagRunType.DATASET_TRIGGERED,
execution_date=exec_date,
data_interval=data_interval,
state=DagRunState.QUEUED,
external_trigger=False,
session=session,
dag_hash=dag_hash,
creating_job_id=self.job.id,
)
Stats.incr("dataset.triggered_dagruns")
dag_run.consumed_dataset_events.extend(dataset_events)
if dataset_events:
dag_run = dag.create_dagrun(
run_id=run_id,
run_type=DagRunType.DATASET_TRIGGERED,
execution_date=exec_date,
data_interval=data_interval,
state=DagRunState.QUEUED,
external_trigger=False,
session=session,
dag_hash=dag_hash,
creating_job_id=self.job.id,
)
Stats.incr("dataset.triggered_dagruns")
dag_run.consumed_dataset_events.extend(dataset_events)
session.execute(
delete(DatasetDagRunQueue).where(DatasetDagRunQueue.target_dag_id == dag_run.dag_id)
delete(DatasetDagRunQueue).where(
DatasetDagRunQueue.target_dag_id == dag.dag_id,
DatasetDagRunQueue.created_at <= exec_date,

@kaxil kaxil Mar 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The created_at <= exec_date scoping correctly preserves newer DDRQ rows that arrive during processing. But when dataset_events is empty and DDRQ rows are deleted without creating a DagRun, there's no log message. Operators debugging "my dataset-triggered DAG didn't fire" would have zero visibility into this silent cleanup. A self.log.warning(...) here would help.

)
)
Comment on lines +1492 to 1511

Copilot AI Apr 2, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dag_run is now created only when dataset_events is non-empty, but the subsequent DELETE no longer depends on dag_run. This fixes the previous unbound reference, but it introduces a more serious logic issue: when dataset_events is empty, the code still deletes from DatasetDagRunQueue for the DAG where created_at <= exec_date. That can clear readiness queue entries without creating a DagRun, potentially dropping triggers and preventing future scheduling. Consider moving the DELETE (or at least the created_at <= exec_date cleanup) inside the if dataset_events: block, or otherwise ensure the queue is only pruned when a run is actually created/consumed.

Copilot uses AI. Check for mistakes.

def _should_update_dag_next_dagruns(
Expand Down
94 changes: 94 additions & 0 deletions tests/datasets/test_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,3 +228,97 @@ def test_create_datasets_notifies_dataset_listener(self, session):
# Ensure the listener was notified
assert len(dataset_listener.created) == 1
assert dataset_listener.created[0].uri == dsm.uri

@pytest.mark.skip_if_database_isolation_mode
def test_slow_path_queue_dagruns_updates_timestamp(self, session):
"""
On non-Postgres backends, _slow_path_queue_dagruns should update created_at
when a DatasetDagRunQueue record already exists for the same (dataset_id, target_dag_id).
"""
from datetime import timedelta

from airflow.utils import timezone

dsem = DatasetManager()
ds = DatasetModel(uri="test_slow_path_update_ts")
dag = DagModel(dag_id="test_slow_path_dag", is_active=True)
session.add_all([ds, dag])
session.flush()

old_ts = timezone.utcnow() - timedelta(days=1)
session.add(DatasetDagRunQueue(dataset_id=ds.id, target_dag_id=dag.dag_id, created_at=old_ts))
session.commit()

before_update = timezone.utcnow()
dsem._slow_path_queue_dagruns(dataset_id=ds.id, dags_to_queue={dag}, session=session)
session.commit()

updated = (
session.query(DatasetDagRunQueue)
.filter_by(dataset_id=ds.id, target_dag_id=dag.dag_id)
.one()
)
assert updated.created_at >= before_update

@pytest.mark.skip_if_database_isolation_mode
@pytest.mark.backend("postgres")
def test_postgres_queue_dagruns_updates_timestamp(self, session):
"""
On PostgreSQL, _postgres_queue_dagruns should update created_at via ON CONFLICT DO UPDATE
when a DatasetDagRunQueue record already exists for the same (dataset_id, target_dag_id).
"""
from datetime import timedelta

from airflow.utils import timezone

dsem = DatasetManager()
ds = DatasetModel(uri="test_postgres_update_ts")
dag = DagModel(dag_id="test_postgres_dag", is_active=True)
session.add_all([ds, dag])
session.flush()

old_ts = timezone.utcnow() - timedelta(days=1)
session.add(DatasetDagRunQueue(dataset_id=ds.id, target_dag_id=dag.dag_id, created_at=old_ts))
session.commit()

before_update = timezone.utcnow()
dsem._postgres_queue_dagruns(dataset_id=ds.id, dags_to_queue={dag}, session=session)
session.commit()

updated = (
session.query(DatasetDagRunQueue)
.filter_by(dataset_id=ds.id, target_dag_id=dag.dag_id)
.one()
)
assert updated.created_at >= before_update

@pytest.mark.skip_if_database_isolation_mode
@pytest.mark.backend("postgres")
def test_postgres_queue_dagruns_where_guard_prevents_backwards_drift(self, session):
"""
The WHERE guard on the Postgres upsert should prevent a slower transaction
from overwriting a newer created_at with an older one.
"""
from datetime import timedelta

from airflow.utils import timezone

dsem = DatasetManager()
ds = DatasetModel(uri="test_postgres_where_guard")
dag = DagModel(dag_id="test_pg_guard_dag", is_active=True)
session.add_all([ds, dag])
session.flush()

newer_ts = timezone.utcnow()
session.add(DatasetDagRunQueue(dataset_id=ds.id, target_dag_id=dag.dag_id, created_at=newer_ts))
session.commit()

dsem._postgres_queue_dagruns(dataset_id=ds.id, dags_to_queue={dag}, session=session)
session.commit()
Comment on lines +297 to +317

Copilot AI Apr 2, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

test_postgres_queue_dagruns_where_guard_prevents_backwards_drift doesn’t currently validate the WHERE-guard behavior it describes. The test sets an existing created_at to newer_ts, then calls _postgres_queue_dagruns(), but that method always uses timezone.utcnow() for the upsert value—i.e. a timestamp newer than newer_ts—so the update is expected to succeed regardless of the WHERE (existing < excluded) guard. To actually regression-test “backwards drift prevention”, freeze/patch timezone.utcnow() (e.g. via the repo’s time_machine fixture) to return an older timestamp for the upsert, then assert created_at remains newer_ts (or unchanged).

Copilot uses AI. Check for mistakes.

updated = (
session.query(DatasetDagRunQueue)
.filter_by(dataset_id=ds.id, target_dag_id=dag.dag_id)
.one()
)
assert updated.created_at >= newer_ts
Loading