You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ph...@apache.org on 2024/07/25 02:13:25 UTC
(airflow) branch main updated: Extend dataset dependencies (#40868)
This is an automated email from the ASF dual-hosted git repository.
phanikumv 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 39971f86d7 Extend dataset dependencies (#40868)
39971f86d7 is described below
commit 39971f86d7f3bb33a0e2f1b4d78232f415085845
Author: Wei Lee <we...@gmail.com>
AuthorDate: Thu Jul 25 10:13:18 2024 +0800
Extend dataset dependencies (#40868)
---
airflow/cli/commands/dag_command.py | 6 +-
airflow/datasets/__init__.py | 81 ++++++++++++++++++++++++++-
airflow/models/serialized_dag.py | 3 +-
airflow/serialization/dag_dependency.py | 43 ++++++++++++++
airflow/serialization/serialized_objects.py | 54 ++++++------------
airflow/timetables/base.py | 4 ++
airflow/utils/dot_renderer.py | 2 +-
airflow/www/views.py | 21 +++++--
tests/serialization/test_dag_serialization.py | 26 ++++++++-
tests/utils/test_dot_renderer.py | 2 +-
10 files changed, 195 insertions(+), 47 deletions(-)
diff --git a/airflow/cli/commands/dag_command.py b/airflow/cli/commands/dag_command.py
index c574d725e2..dac61d0da5 100644
--- a/airflow/cli/commands/dag_command.py
+++ b/airflow/cli/commands/dag_command.py
@@ -265,7 +265,11 @@ def set_is_paused(is_paused: bool, args) -> None:
@providers_configuration_loaded
def dag_dependencies_show(args) -> None:
"""Display DAG dependencies, save to file or show as imgcat image."""
- dot = render_dag_dependencies(SerializedDagModel.get_dag_dependencies())
+ deduplicated_dag_dependencies = {
+ dag_id: list(set(dag_dependencies))
+ for dag_id, dag_dependencies in SerializedDagModel.get_dag_dependencies().items()
+ }
+ dot = render_dag_dependencies(deduplicated_dag_dependencies)
filename = args.save
imgcat = args.imgcat
diff --git a/airflow/datasets/__init__.py b/airflow/datasets/__init__.py
index dd27338189..873d314136 100644
--- a/airflow/datasets/__init__.py
+++ b/airflow/datasets/__init__.py
@@ -20,11 +20,12 @@ from __future__ import annotations
import os
import urllib.parse
import warnings
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Iterable, Iterator
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Iterable, Iterator, cast
import attr
from sqlalchemy import select
+from airflow.serialization.dag_dependency import DagDependency
from airflow.typing_compat import TypedDict
from airflow.utils.session import NEW_SESSION, provide_session
@@ -191,6 +192,14 @@ class BaseDataset:
def iter_datasets(self) -> Iterator[tuple[str, Dataset]]:
raise NotImplementedError
+ def iter_dag_dependencies(self, *, source: str, target: str) -> Iterator[DagDependency]:
+ """
+ Iterate a base dataset as dag dependency.
+
+ :meta private:
+ """
+ raise NotImplementedError
+
@attr.define()
class DatasetAlias(BaseDataset):
@@ -206,6 +215,19 @@ class DatasetAlias(BaseDataset):
def __hash__(self) -> int:
return hash(self.name)
+ def iter_dag_dependencies(self, *, source: str, target: str) -> Iterator[DagDependency]:
+ """
+ Iterate a dataset alias as dag dependency.
+
+ :meta private:
+ """
+ yield DagDependency(
+ source=source or "dataset-alias",
+ target=target or "dataset-alias",
+ dependency_type="dataset-alias",
+ dependency_id=self.name,
+ )
+
class DatasetAliasEvent(TypedDict):
"""A represeation of dataset event to be triggered by a dataset alias."""
@@ -273,6 +295,19 @@ class Dataset(os.PathLike, BaseDataset):
def evaluate(self, statuses: dict[str, bool]) -> bool:
return statuses.get(self.uri, False)
+ def iter_dag_dependencies(self, *, source: str, target: str) -> Iterator[DagDependency]:
+ """
+ Iterate a dataset as dag dependency.
+
+ :meta private:
+ """
+ yield DagDependency(
+ source=source or "dataset",
+ target=target or "dataset",
+ dependency_type="dataset",
+ dependency_id=self.uri,
+ )
+
class _DatasetBooleanCondition(BaseDataset):
"""Base class for dataset boolean logic."""
@@ -299,6 +334,15 @@ class _DatasetBooleanCondition(BaseDataset):
yield k, v
seen.add(k)
+ def iter_dag_dependencies(self, *, source: str, target: str) -> Iterator[DagDependency]:
+ """
+ Iterate dataset, dataset aliases and their resolved datasets as dag dependency.
+
+ :meta private:
+ """
+ for obj in self.objects:
+ yield from obj.iter_dag_dependencies(source=source, target=target)
+
class DatasetAny(_DatasetBooleanCondition):
"""Use to combine datasets schedule references in an "and" relationship."""
@@ -334,6 +378,9 @@ class _DatasetAliasCondition(DatasetAny):
self.name = name
self.objects = expand_alias_to_datasets(name)
+ def __repr__(self) -> str:
+ return f"_DatasetAliasCondition({', '.join(map(str, self.objects))})"
+
def as_expression(self) -> Any:
"""
Serialize the dataset into its scheduling expression.
@@ -342,6 +389,38 @@ class _DatasetAliasCondition(DatasetAny):
"""
return {"alias": self.name}
+ def iter_dag_dependencies(self, *, source: str = "", target: str = "") -> Iterator[DagDependency]:
+ """
+ Iterate a dataset alias and its resolved datasets as dag dependency.
+
+ :meta private:
+ """
+ if self.objects:
+ for obj in self.objects:
+ dataset = cast(Dataset, obj)
+ uri = dataset.uri
+ # dataset
+ yield DagDependency(
+ source=f"dataset-alias:{self.name}" if source else "dataset",
+ target="dataset" if source else f"dataset-alias:{self.name}",
+ dependency_type="dataset",
+ dependency_id=uri,
+ )
+ # dataset alias
+ yield DagDependency(
+ source=source or f"dataset:{uri}",
+ target=target or f"dataset:{uri}",
+ dependency_type="dataset-alias",
+ dependency_id=self.name,
+ )
+ else:
+ yield DagDependency(
+ source=source or "dataset-alias",
+ target=target or "dataset-alias",
+ dependency_type="dataset-alias",
+ dependency_id=self.name,
+ )
+
class DatasetAll(_DatasetBooleanCondition):
"""Use to combine datasets schedule references in an "or" relationship."""
diff --git a/airflow/models/serialized_dag.py b/airflow/models/serialized_dag.py
index 646e1c61f1..99495f81b2 100644
--- a/airflow/models/serialized_dag.py
+++ b/airflow/models/serialized_dag.py
@@ -35,7 +35,8 @@ from airflow.models.base import ID_LEN, Base
from airflow.models.dag import DagModel
from airflow.models.dagcode import DagCode
from airflow.models.dagrun import DagRun
-from airflow.serialization.serialized_objects import DagDependency, SerializedDAG
+from airflow.serialization.dag_dependency import DagDependency
+from airflow.serialization.serialized_objects import SerializedDAG
from airflow.settings import COMPRESS_SERIALIZED_DAGS, MIN_SERIALIZED_DAG_UPDATE_INTERVAL, json
from airflow.utils import timezone
from airflow.utils.hashlib_wrapper import md5
diff --git a/airflow/serialization/dag_dependency.py b/airflow/serialization/dag_dependency.py
new file mode 100644
index 0000000000..bff1b39ebe
--- /dev/null
+++ b/airflow/serialization/dag_dependency.py
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, order=True)
+class DagDependency:
+ """
+ Dataclass for representing dependencies between DAGs.
+
+ These are calculated during serialization and attached to serialized DAGs.
+ """
+
+ source: str
+ target: str
+ dependency_type: str
+ dependency_id: str | None = None
+
+ @property
+ def node_id(self):
+ """Node ID for graph rendering."""
+ val = f"{self.dependency_type}"
+ if self.dependency_type not in ("dataset", "dataset-alias"):
+ val += f":{self.source}:{self.target}"
+ if self.dependency_id:
+ val += f":{self.dependency_id}"
+ return val
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
index 90797d9535..4d6a5b7e57 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -25,7 +25,6 @@ import inspect
import logging
import warnings
import weakref
-from dataclasses import dataclass
from inspect import signature
from textwrap import dedent
from typing import TYPE_CHECKING, Any, Collection, Iterable, Mapping, NamedTuple, Union
@@ -38,7 +37,14 @@ from pendulum.tz.timezone import FixedTimezone, Timezone
from airflow.callbacks.callback_requests import DagCallbackRequest, SlaCallbackRequest, TaskCallbackRequest
from airflow.compat.functools import cache
from airflow.configuration import conf
-from airflow.datasets import BaseDataset, Dataset, DatasetAlias, DatasetAll, DatasetAny
+from airflow.datasets import (
+ BaseDataset,
+ Dataset,
+ DatasetAlias,
+ DatasetAll,
+ DatasetAny,
+ _DatasetAliasCondition,
+)
from airflow.exceptions import AirflowException, RemovedInAirflow3Warning, SerializationError, TaskDeferred
from airflow.jobs.job import Job
from airflow.models.baseoperator import BaseOperator
@@ -52,6 +58,7 @@ from airflow.models.taskinstance import SimpleTaskInstance, TaskInstance
from airflow.models.tasklog import LogTemplate
from airflow.models.xcom_arg import XComArg, deserialize_xcom_arg, serialize_xcom_arg
from airflow.providers_manager import ProvidersManager
+from airflow.serialization.dag_dependency import DagDependency
from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding
from airflow.serialization.helpers import serialize_template_field
from airflow.serialization.json_schema import load_dag_schema
@@ -966,6 +973,10 @@ class DependencyDetector:
dependency_id=obj.uri,
)
)
+ elif isinstance(obj, DatasetAlias):
+ cond = _DatasetAliasCondition(obj.name)
+
+ deps.extend(cond.iter_dag_dependencies(source=task.dag_id, target=""))
return deps
@staticmethod
@@ -973,13 +984,8 @@ class DependencyDetector:
"""Detect dependencies set directly on the DAG object."""
if not dag:
return
- for uri, _ in dag.timetable.dataset_condition.iter_datasets():
- yield DagDependency(
- source="dataset",
- target=dag.dag_id,
- dependency_type="dataset",
- dependency_id=uri,
- )
+
+ yield from dag.timetable.dataset_condition.iter_dag_dependencies(source="", target=dag.dag_id)
class SerializedBaseOperator(BaseOperator, BaseSerialization):
@@ -1562,12 +1568,12 @@ class SerializedDAG(DAG, BaseSerialization):
serialized_dag["tasks"] = [cls.serialize(task) for _, task in dag.task_dict.items()]
- dag_deps = {
+ dag_deps = [
dep
for task in dag.task_dict.values()
for dep in SerializedBaseOperator.detect_dependencies(task)
- }
- dag_deps.update(DependencyDetector.detect_dag_dependencies(dag))
+ ]
+ dag_deps.extend(DependencyDetector.detect_dag_dependencies(dag))
serialized_dag["dag_dependencies"] = [x.__dict__ for x in sorted(dag_deps)]
serialized_dag["_task_group"] = TaskGroupSerialization.serialize_task_group(dag.task_group)
@@ -1773,30 +1779,6 @@ class TaskGroupSerialization(BaseSerialization):
return group
-@dataclass(frozen=True, order=True)
-class DagDependency:
- """
- Dataclass for representing dependencies between DAGs.
-
- These are calculated during serialization and attached to serialized DAGs.
- """
-
- source: str
- target: str
- dependency_type: str
- dependency_id: str | None = None
-
- @property
- def node_id(self):
- """Node ID for graph rendering."""
- val = f"{self.dependency_type}"
- if self.dependency_type != "dataset":
- val += f":{self.source}:{self.target}"
- if self.dependency_id:
- val += f":{self.dependency_id}"
- return val
-
-
def _has_kubernetes() -> bool:
global HAS_KUBERNETES
if "HAS_KUBERNETES" in globals():
diff --git a/airflow/timetables/base.py b/airflow/timetables/base.py
index 04a1a2fb4b..9e173656f1 100644
--- a/airflow/timetables/base.py
+++ b/airflow/timetables/base.py
@@ -26,6 +26,7 @@ if TYPE_CHECKING:
from pendulum import DateTime
from airflow.datasets import Dataset
+ from airflow.serialization.dag_dependency import DagDependency
from airflow.utils.types import DagRunType
@@ -57,6 +58,9 @@ class _NullDataset(BaseDataset):
def iter_datasets(self) -> Iterator[tuple[str, Dataset]]:
return iter(())
+ def iter_dag_dependencies(self, source, target) -> Iterator[DagDependency]:
+ return iter(())
+
class DataInterval(NamedTuple):
"""
diff --git a/airflow/utils/dot_renderer.py b/airflow/utils/dot_renderer.py
index 4cc3be5ddf..ec0b836cbc 100644
--- a/airflow/utils/dot_renderer.py
+++ b/airflow/utils/dot_renderer.py
@@ -44,7 +44,7 @@ if TYPE_CHECKING:
from airflow.models import TaskInstance
from airflow.models.dag import DAG
from airflow.models.taskmixin import DependencyMixin
- from airflow.serialization.serialized_objects import DagDependency
+ from airflow.serialization.dag_dependency import DagDependency
def _refine_color(color: str):
diff --git a/airflow/www/views.py b/airflow/www/views.py
index 568b4bfe39..6990890fe0 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -3503,16 +3503,27 @@ class Airflow(AirflowBaseView):
dag_node_id = f"dag:{dag}"
if dag_node_id not in nodes_dict:
for dep in dependencies:
- if dep.dependency_type == "dag" or dep.dependency_type == "dataset":
+ if dep.dependency_type in ("dag", "dataset", "dataset-alias"):
+ # add node
nodes_dict[dag_node_id] = node_dict(dag_node_id, dag, "dag")
if dep.node_id not in nodes_dict:
nodes_dict[dep.node_id] = node_dict(
dep.node_id, dep.dependency_id, dep.dependency_type
)
- if dep.source != "dataset":
- edge_tuples.add((f"dag:{dep.source}", dep.node_id))
- if dep.target != "dataset":
- edge_tuples.add((dep.node_id, f"dag:{dep.target}"))
+
+ # add edge
+
+ # not start dep
+ if dep.source != dep.dependency_type:
+ source = dep.source if ":" in dep.source else f"dag:{dep.source}"
+ target = dep.node_id
+ edge_tuples.add((source, target))
+
+ # not end dep
+ if dep.target != dep.dependency_type:
+ source = dep.node_id
+ target = dep.target if ":" in dep.target else f"dag:{dep.target}"
+ edge_tuples.add((source, target))
nodes = list(nodes_dict.values())
edges = [{"source": source, "target": target} for source, target in edge_tuples]
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
index 65c4650c7a..8650378854 100644
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -60,11 +60,11 @@ from airflow.operators.empty import EmptyOperator
from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator
from airflow.security import permissions
from airflow.sensors.bash import BashSensor
+from airflow.serialization.dag_dependency import DagDependency
from airflow.serialization.enums import Encoding
from airflow.serialization.json_schema import load_dag_schema_dict
from airflow.serialization.serialized_objects import (
BaseSerialization,
- DagDependency,
DependencyDetector,
SerializedBaseOperator,
SerializedDAG,
@@ -1796,6 +1796,30 @@ class TestStringifiedDAGs:
"dependency_type": "dataset",
"dependency_id": "d1",
},
+ {
+ "dependency_id": "d1",
+ "dependency_type": "dataset",
+ "source": "dataset",
+ "target": "test",
+ },
+ {
+ "dependency_id": "d1",
+ "dependency_type": "dataset",
+ "source": "dataset",
+ "target": "test",
+ },
+ {
+ "dependency_id": "d1",
+ "dependency_type": "dataset",
+ "source": "dataset",
+ "target": "test",
+ },
+ {
+ "dependency_id": "d1",
+ "dependency_type": "dataset",
+ "source": "dataset",
+ "target": "test",
+ },
],
key=lambda x: tuple(x.values()),
)
diff --git a/tests/utils/test_dot_renderer.py b/tests/utils/test_dot_renderer.py
index 4c57f8026e..2408a67fe2 100644
--- a/tests/utils/test_dot_renderer.py
+++ b/tests/utils/test_dot_renderer.py
@@ -26,7 +26,7 @@ from airflow.models.dag import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.empty import EmptyOperator
from airflow.operators.python import PythonOperator
-from airflow.serialization.serialized_objects import DagDependency
+from airflow.serialization.dag_dependency import DagDependency
from airflow.utils import dot_renderer, timezone
from airflow.utils.state import State
from airflow.utils.task_group import TaskGroup