You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2022/08/17 21:32:26 UTC
[airflow] branch main updated: Used many-to-many relation for finding attached dataset events. (#25758)
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f3ca4f505c Used many-to-many relation for finding attached dataset events. (#25758)
f3ca4f505c is described below
commit f3ca4f505c63a49496451b9eb4dd6df397bf7f6b
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Wed Aug 17 22:32:11 2022 +0100
Used many-to-many relation for finding attached dataset events. (#25758)
The previous PR added this join table -- by using it we make the
association of which events created a DagRun fixed at DagRun creation
time; previously if you deleted old DagRuns then the oldest remaining
run would "collect" all the oldest events.
---
.../api_connexion/endpoints/dag_run_endpoint.py | 34 +----
airflow/jobs/scheduler_job.py | 2 +-
airflow/models/dataset.py | 2 +-
.../endpoints/test_dag_run_endpoint.py | 152 +++++----------------
tests/jobs/test_scheduler_job.py | 2 +-
5 files changed, 37 insertions(+), 155 deletions(-)
diff --git a/airflow/api_connexion/endpoints/dag_run_endpoint.py b/airflow/api_connexion/endpoints/dag_run_endpoint.py
index 30e01eedec..eb7026b5df 100644
--- a/airflow/api_connexion/endpoints/dag_run_endpoint.py
+++ b/airflow/api_connexion/endpoints/dag_run_endpoint.py
@@ -51,7 +51,6 @@ from airflow.api_connexion.schemas.task_instance_schema import (
)
from airflow.api_connexion.types import APIResponse
from airflow.models import DagModel, DagRun
-from airflow.models.dataset import DatasetDagRef, DatasetEvent
from airflow.security import permissions
from airflow.utils.airflow_flask_app import get_airflow_app
from airflow.utils.session import NEW_SESSION, provide_session
@@ -116,43 +115,12 @@ def get_upstream_dataset_events(
"DAGRun not found",
detail=f"DAGRun with DAG ID: '{dag_id}' and DagRun ID: '{dag_run_id}' not found",
)
- events = _get_upstream_dataset_events(dag_run=dag_run, session=session)
+ events = dag_run.consumed_dataset_events
return dataset_event_collection_schema.dump(
DatasetEventCollection(dataset_events=events, total_entries=len(events))
)
-def _get_upstream_dataset_events(*, dag_run: DagRun, session: Session) -> List["DagRun"]:
- if not dag_run.run_type == DagRunType.DATASET_TRIGGERED:
- return []
-
- previous_dag_run = (
- session.query(DagRun)
- .filter(
- DagRun.dag_id == dag_run.dag_id,
- DagRun.execution_date < dag_run.execution_date,
- DagRun.run_type == DagRunType.DATASET_TRIGGERED,
- )
- .order_by(DagRun.execution_date.desc())
- .first()
- )
-
- dataset_event_filters = [
- DatasetDagRef.dag_id == dag_run.dag_id,
- DatasetEvent.timestamp <= dag_run.execution_date,
- ]
- if previous_dag_run:
- dataset_event_filters.append(DatasetEvent.timestamp > previous_dag_run.execution_date)
- dataset_events = (
- session.query(DatasetEvent)
- .join(DatasetDagRef, DatasetEvent.dataset_id == DatasetDagRef.dataset_id)
- .filter(*dataset_event_filters)
- .order_by(DatasetEvent.timestamp)
- .all()
- )
- return dataset_events
-
-
def _fetch_dag_runs(
query: Query,
*,
diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job.py
index 571bed18b2..86973fc9e4 100644
--- a/airflow/jobs/scheduler_job.py
+++ b/airflow/jobs/scheduler_job.py
@@ -1130,7 +1130,7 @@ class SchedulerJob(BaseJob):
.all()
)
- dag_run.dataset_events.extend(dataset_events)
+ dag_run.consumed_dataset_events.extend(dataset_events)
session.query(DatasetDagRunQueue).filter(
DatasetDagRunQueue.target_dag_id == dag_run.dag_id
).delete()
diff --git a/airflow/models/dataset.py b/airflow/models/dataset.py
index c50266e07a..f2ee80dac7 100644
--- a/airflow/models/dataset.py
+++ b/airflow/models/dataset.py
@@ -271,7 +271,7 @@ class DatasetEvent(Base):
created_dagruns = relationship(
"DagRun",
secondary=association_table,
- backref="dataset_events",
+ backref="consumed_dataset_events",
)
source_task_instance = relationship(
diff --git a/tests/api_connexion/endpoints/test_dag_run_endpoint.py b/tests/api_connexion/endpoints/test_dag_run_endpoint.py
index 23c5549db4..1a7c129854 100644
--- a/tests/api_connexion/endpoints/test_dag_run_endpoint.py
+++ b/tests/api_connexion/endpoints/test_dag_run_endpoint.py
@@ -17,7 +17,6 @@
from datetime import timedelta
from unittest import mock
-import pendulum
import pytest
from freezegun import freeze_time
from parameterized import parameterized
@@ -30,7 +29,7 @@ from airflow.operators.empty import EmptyOperator
from airflow.security import permissions
from airflow.utils import timezone
from airflow.utils.session import create_session, provide_session
-from airflow.utils.state import DagRunState, State
+from airflow.utils.state import State
from airflow.utils.types import DagRunType
from tests.test_utils.api_connexion_utils import assert_401, create_user, delete_roles, delete_user
from tests.test_utils.config import conf_vars
@@ -1490,118 +1489,33 @@ class TestClearDagRun(TestDagRunEndpoint):
@pytest.mark.need_serialized_dag
-def test__get_upstream_dataset_events_no_prior(configured_app, dag_maker):
- """If no prior dag runs, return all events"""
- from airflow.api_connexion.endpoints.dag_run_endpoint import _get_upstream_dataset_events
-
- dataset1a = Dataset(uri="ds1a")
- dataset1b = Dataset(uri="ds1b")
- # setup dags and datasets
- with dag_maker(dag_id="datasets-2", schedule=[dataset1a, dataset1b]):
- pass
-
- # add 5 events
- session = dag_maker.session
- ds1a_id = session.query(DatasetModel.id).filter_by(uri=dataset1a.uri).scalar()
- ds1b_id = session.query(DatasetModel.id).filter_by(uri=dataset1b.uri).scalar()
- session.add_all([DatasetEvent(dataset_id=ds1a_id), DatasetEvent(dataset_id=ds1b_id)])
- session.add_all([DatasetEvent(dataset_id=ds1a_id), DatasetEvent(dataset_id=ds1b_id)])
- session.add_all([DatasetEvent(dataset_id=ds1a_id)])
- session.flush()
-
- # create a single dag run, no prior dag runs
- dr = dag_maker.create_dagrun(
- run_id="run", run_type=DagRunType.DATASET_TRIGGERED, execution_date=timezone.utcnow()
- )
-
- # check result
- events = _get_upstream_dataset_events(dag_run=dr, session=session)
- assert len(events) == 5
-
-
-@pytest.mark.need_serialized_dag
-def test__get_upstream_dataset_events_with_prior(configured_app, dag_maker):
- """
- Events returned should be those that occurred after last DATASET_TRIGGERED
- dag run and up to the exec date of current dag run.
- """
- from airflow.api_connexion.endpoints.dag_run_endpoint import _get_upstream_dataset_events
-
- # setup dags and datasets
- dataset1a = Dataset(uri="ds1a")
- dataset1b = Dataset(uri="ds1b")
- # setup dags and datasets
- with dag_maker(dag_id="datasets-2", schedule=[dataset1a, dataset1b]):
- pass
-
- session = dag_maker.session
- ds1a_id = session.query(DatasetModel.id).filter_by(uri=dataset1a.uri).scalar()
- ds1b_id = session.query(DatasetModel.id).filter_by(uri=dataset1b.uri).scalar()
-
- # add 2 events, then a dag run, then 3 events, then another dag run then another event
- first_timestamp = pendulum.datetime(2022, 1, 1, tz='UTC')
- session.add_all(
- [
- DatasetEvent(dataset_id=ds1a_id, timestamp=first_timestamp),
- DatasetEvent(dataset_id=ds1b_id, timestamp=first_timestamp),
- ]
- )
- dag_maker.create_dagrun(
- run_id='run-1',
- run_type=DagRunType.DATASET_TRIGGERED,
- execution_date=first_timestamp.add(microseconds=1000),
- )
- session.add_all(
- [
- DatasetEvent(dataset_id=ds1a_id, timestamp=first_timestamp.add(microseconds=2000)),
- DatasetEvent(dataset_id=ds1b_id, timestamp=first_timestamp.add(microseconds=3000)),
- DatasetEvent(dataset_id=ds1b_id, timestamp=first_timestamp.add(microseconds=4000)),
- ]
- )
- dag_maker.create_dagrun(
- run_id='run-2',
- run_type=DagRunType.MANUAL,
- execution_date=first_timestamp.add(microseconds=3000),
- )
- dr3 = dag_maker.create_dagrun(
- run_id='run-3',
- run_type=DagRunType.DATASET_TRIGGERED,
- execution_date=first_timestamp.add(microseconds=4000), # exact same time as 3rd event in window
- )
- session.add_all([DatasetEvent(dataset_id=ds1a_id, timestamp=first_timestamp.add(microseconds=5000))])
- session.flush()
-
- events = _get_upstream_dataset_events(dag_run=dr3, session=session)
-
- event_times = [x.timestamp for x in events]
- assert event_times == [
- first_timestamp.add(microseconds=2000),
- first_timestamp.add(microseconds=3000),
- first_timestamp.add(microseconds=4000),
- ]
-
-
class TestGetDagRunDatasetTriggerEvents(TestDagRunEndpoint):
- @mock.patch('airflow.api_connexion.endpoints.dag_run_endpoint._get_upstream_dataset_events')
- def test_should_respond_200(self, mock_get_events, session):
- dagrun_model = DagRun(
- dag_id="TEST_DAG_ID",
- run_id="TEST_DAG_RUN_ID",
- run_type=DagRunType.DATASET_TRIGGERED,
- execution_date=timezone.parse(self.default_time),
- start_date=timezone.parse(self.default_time),
- external_trigger=True,
- state=DagRunState.RUNNING,
+ def test_should_respond_200(self, dag_maker, session):
+ dataset1 = Dataset(uri="ds1")
+
+ with dag_maker(dag_id="source_dag", start_date=timezone.utcnow(), session=session):
+ EmptyOperator(task_id="task", outlets=[dataset1])
+ dr = dag_maker.create_dagrun()
+ ti = dr.task_instances[0]
+
+ ds1_id = session.query(DatasetModel.id).filter_by(uri=dataset1.uri).scalar()
+ event = DatasetEvent(
+ dataset_id=ds1_id,
+ source_task_id=ti.task_id,
+ source_dag_id=ti.dag_id,
+ source_run_id=ti.run_id,
+ source_map_index=ti.map_index,
)
- session.add(dagrun_model)
+ session.add(event)
+
+ with dag_maker(dag_id="TEST_DAG_ID", start_date=timezone.utcnow(), session=session):
+ pass
+ dr = dag_maker.create_dagrun(run_id="TEST_DAG_RUN_ID", run_type=DagRunType.DATASET_TRIGGERED)
+ dr.consumed_dataset_events.append(event)
+
session.commit()
- result = session.query(DagRun).all()
- assert len(result) == 1
- created_at = pendulum.now('UTC')
- # make sure whatever is returned by this func is what comes out in response.
- d = DatasetEvent(dataset_id=1, timestamp=created_at)
- d.dataset = DatasetModel(id=1, uri='hello', created_at=created_at, updated_at=created_at)
- mock_get_events.return_value = [d]
+ assert event.timestamp
+
response = self.client.get(
"api/v1/dags/TEST_DAG_ID/dagRuns/TEST_DAG_RUN_ID/upstreamDatasetEvents",
environ_overrides={'REMOTE_USER': "test"},
@@ -1610,15 +1524,15 @@ class TestGetDagRunDatasetTriggerEvents(TestDagRunEndpoint):
expected_response = {
'dataset_events': [
{
- 'timestamp': str(created_at),
- 'dataset_id': 1,
- 'dataset_uri': d.dataset.uri,
+ 'timestamp': event.timestamp.isoformat(),
+ 'dataset_id': ds1_id,
+ 'dataset_uri': dataset1.uri,
'extra': {},
- 'id': None,
- 'source_dag_id': None,
- 'source_map_index': None,
- 'source_run_id': None,
- 'source_task_id': None,
+ 'id': event.id,
+ 'source_dag_id': ti.dag_id,
+ 'source_map_index': ti.map_index,
+ 'source_run_id': ti.run_id,
+ 'source_task_id': ti.task_id,
}
],
'total_entries': 1,
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 21bdd16625..1792fc6241 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -3083,7 +3083,7 @@ class TestSchedulerJob:
created_run = session.query(DagRun).filter(DagRun.dag_id == dag3.dag_id).one()
assert created_run.state == State.QUEUED
assert created_run.start_date is None
- assert created_run.dataset_events == [event]
+ assert created_run.consumed_dataset_events == [event]
# dag3 DDRQ record should still be there since the dag run was *not* triggered
assert session.query(DatasetDagRunQueue).filter(DagRun.dag_id == dag3.dag_id).one() is not None
# dag2 should not be triggered since it depends on both dataset 1 and 2