Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 4 additions & 0 deletions airflow/jobs/scheduler_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -1280,6 +1280,10 @@ def _create_dag_runs_dataset_triggered(
]
if previous_dag_run:
dataset_event_filters.append(DatasetEvent.timestamp > previous_dag_run.execution_date)
else:
dataset_event_filters.append(
DatasetEvent.timestamp > DagScheduleDatasetReference.created_at
)

dataset_events = session.scalars(
select(DatasetEvent)
Expand Down
8 changes: 8 additions & 0 deletions newsfragments/39603.significant.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
New DAGs will no longer receive historical dataset events

In the past, new DAGs received historical dataset events from the
very first dataset event in the system. This has been corrected,
and now new DAGs will only receive dataset events that have occurred
after their creation. Although this change may disrupt existing workflows,
the previous behavior is considered a bug, and this update ensures more
accurate and expected DAG execution.
120 changes: 111 additions & 9 deletions tests/jobs/test_scheduler_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -3793,6 +3793,10 @@ def test_create_dag_runs_datasets(self, session, dag_maker):
dataset1 = Dataset(uri="ds1")
dataset2 = Dataset(uri="ds2")

with dag_maker(dag_id="datasets-consumer-single", schedule=[dataset1]):
Comment thread
potiuk marked this conversation as resolved.
pass
dag2 = dag_maker.dag

with dag_maker(dag_id="datasets-1", start_date=timezone.utcnow(), session=session):
BashOperator(task_id="task", bash_command="echo 1", outlets=[dataset1])
dr = dag_maker.create_dagrun(
Expand Down Expand Up @@ -3830,9 +3834,6 @@ def test_create_dag_runs_datasets(self, session, dag_maker):

with dag_maker(dag_id="datasets-consumer-multiple", schedule=[dataset1, dataset2]):
pass
dag2 = dag_maker.dag
with dag_maker(dag_id="datasets-consumer-single", schedule=[dataset1]):
pass
dag3 = dag_maker.dag

session = dag_maker.session
Expand All @@ -3856,8 +3857,8 @@ def dict_from_obj(obj):
"""Get dict of column attrs from SqlAlchemy object."""
return {k.key: obj.__dict__.get(k) for k in obj.__mapper__.column_attrs}

# dag3 should be triggered since it only depends on dataset1, and it's been queued
created_run = session.query(DagRun).filter(DagRun.dag_id == dag3.dag_id).one()
# dag2 should be triggered since it only depends on dataset1, it's been queued and dataset events landed after DAG was created.
created_run = session.query(DagRun).filter(DagRun.dag_id == dag2.dag_id).one()
assert created_run.state == State.QUEUED
assert created_run.start_date is None

Expand All @@ -3869,13 +3870,114 @@ def dict_from_obj(obj):
assert created_run.data_interval_start == DEFAULT_DATE + timedelta(days=5)
assert created_run.data_interval_end == DEFAULT_DATE + timedelta(days=11)
# dag2 DDRQ record should still be there since the dag run was *not* triggered
assert session.query(DatasetDagRunQueue).filter_by(target_dag_id=dag2.dag_id).one() is not None
assert session.query(DatasetDagRunQueue).filter_by(target_dag_id=dag3.dag_id).one() is not None
# dag2 should not be triggered since it depends on both dataset 1 and 2
assert session.query(DagRun).filter(DagRun.dag_id == dag2.dag_id).one_or_none() is None
assert session.query(DagRun).filter(DagRun.dag_id == dag3.dag_id).one_or_none() is None
# dag3 DDRQ record should be deleted since the dag run was triggered
assert session.query(DatasetDagRunQueue).filter_by(target_dag_id=dag3.dag_id).one_or_none() is None
assert session.query(DatasetDagRunQueue).filter_by(target_dag_id=dag2.dag_id).one_or_none() is None

assert dag2.get_last_dagrun().creating_job_id == scheduler_job.id
Comment thread
potiuk marked this conversation as resolved.

@pytest.mark.need_serialized_dag
def test_new_dagrun_ignores_old_dataset_events(self, session, dag_maker):
"""
Test various invariants of _create_dag_runs.

- That the new DAG should not get dataset events which has timestamp with before dag creation date.
- That the run created is on QUEUED State
- That dag_model has next_dagrun
"""

dataset = Dataset(uri="ds")

with dag_maker(dag_id="datasets-1", start_date=timezone.utcnow(), session=session):
BashOperator(task_id="task", bash_command="echo 1", outlets=[dataset])
dr = dag_maker.create_dagrun(
run_id="run1",
execution_date=(DEFAULT_DATE + timedelta(days=100)),
data_interval=(DEFAULT_DATE + timedelta(days=10), DEFAULT_DATE + timedelta(days=11)),
)

ds_id = session.query(DatasetModel.id).filter_by(uri=dataset.uri).scalar()

event1 = DatasetEvent(
dataset_id=ds_id,
source_task_id="task",
source_dag_id=dr.dag_id,
source_run_id=dr.run_id,
source_map_index=-1,
)
session.add(event1)

# Create a second event, creation time is more recent, but data interval is older
dr = dag_maker.create_dagrun(
run_id="run2",
execution_date=(DEFAULT_DATE + timedelta(days=101)),
data_interval=(DEFAULT_DATE + timedelta(days=5), DEFAULT_DATE + timedelta(days=6)),
)

event2 = DatasetEvent(
dataset_id=ds_id,
source_task_id="task",
source_dag_id=dr.dag_id,
source_run_id=dr.run_id,
source_map_index=-1,
)
session.add(event2)

# Create DAG after dataset events.
with dag_maker(dag_id="datasets-consumer", schedule=[dataset]):
pass
consumer_dag = dag_maker.dag

with dag_maker(dag_id="datasets-1-new", start_date=timezone.utcnow(), session=session):
BashOperator(task_id="task", bash_command="echo 1", outlets=[dataset])
dr = dag_maker.create_dagrun(
run_id="run3",
execution_date=(DEFAULT_DATE + timedelta(days=101)),
data_interval=(DEFAULT_DATE + timedelta(days=5), DEFAULT_DATE + timedelta(days=6)),
)

event3 = DatasetEvent(
dataset_id=ds_id,
source_task_id="task",
source_dag_id=dr.dag_id,
source_run_id=dr.run_id,
source_map_index=-1,
timestamp=timezone.utcnow(),
)
session.add(event3)

session = dag_maker.session
session.add(DatasetDagRunQueue(dataset_id=ds_id, target_dag_id=consumer_dag.dag_id))
session.flush()

scheduler_job = Job(executor=self.null_exec)
self.job_runner = SchedulerJobRunner(job=scheduler_job)

self.job_runner.processor_agent = mock.MagicMock()

with create_session() as session:
self.job_runner._create_dagruns_for_dags(session, session)

def dict_from_obj(obj):
"""Get dict of column attrs from SqlAlchemy object."""
return {k.key: obj.__dict__.get(k) for k in obj.__mapper__.column_attrs}
Comment thread
tosheer marked this conversation as resolved.

# dag should be triggered since it only depends on dataset1, it's been queued and dataset events landed after DAG was created.
created_run = session.query(DagRun).filter(DagRun.dag_id == consumer_dag.dag_id).one()
assert created_run.state == State.QUEUED

# __eq__ isn't defined on DatasetEvent
assert list(map(dict_from_obj, created_run.consumed_dataset_events)) == [dict_from_obj(event3)]

# dag DDRQ record should be deleted since the dag run was triggered
assert (
session.query(DatasetDagRunQueue).filter_by(target_dag_id=consumer_dag.dag_id).one_or_none()
is None
)

assert dag3.get_last_dagrun().creating_job_id == scheduler_job.id
assert consumer_dag.get_last_dagrun().creating_job_id == scheduler_job.id

@pytest.mark.need_serialized_dag
@pytest.mark.parametrize(
Expand Down