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 2024/10/30 14:31:11 UTC
(airflow) 13/31: fix more stest [skip-ci]
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch task-sdk-first-code
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 444016cc59b635cfa937238b3e5ba1e86d6f0f28
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Thu Oct 24 16:25:41 2024 +0100
fix more stest [skip-ci]
---
.pre-commit-config.yaml | 1 +
airflow/models/dag.py | 25 ++-----
airflow/models/taskmixin.py | 24 +-----
.../src/airflow/sdk/definitions/baseoperator.py | 15 ++--
task_sdk/src/airflow/sdk/definitions/dag.py | 86 +++++++++++++++-------
task_sdk/src/airflow/sdk/definitions/node.py | 7 ++
task_sdk/tests/defintions/test_dag.py | 34 +++++++++
tests/models/test_dag.py | 31 +++-----
8 files changed, 128 insertions(+), 95 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7d9cd6b2292..ad4b2529b86 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1190,6 +1190,7 @@ repos:
^providers/src/airflow/providers/ |
^(providers/)?tests/ |
task_sdk/src/airflow/sdk/definitions/dag.py$ |
+ task_sdk/src/airflow/sdk/definitions/node.py$ |
^dev/.*\.py$ |
^scripts/.*\.py$ |
^docker_tests/.*$ |
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index fc0928ddcc1..d133eb43e68 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -298,7 +298,7 @@ if TYPE_CHECKING:
else:
def dag(dag_id: str = "", **kwargs):
- return task_sdk_dag_decorator(dag_id, __DAG_class=DAG, __warnings_stacklevel_delta=3)
+ return task_sdk_dag_decorator(dag_id, __DAG_class=DAG, __warnings_stacklevel_delta=3, **kwargs)
@functools.total_ordering
@@ -419,15 +419,14 @@ class DAG(TaskSDKDag, LoggingMixin):
partial: bool = False
last_loaded: datetime | None = attrs.field(factory=timezone.utcnow)
- on_success_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None
- on_failure_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None
-
- has_on_success_callback: bool = attrs.field(init=False)
- has_on_failure_callback: bool = attrs.field(init=False)
default_view: str = airflow_conf.get_mandatory_value("webserver", "dag_default_view").lower()
orientation: str = airflow_conf.get_mandatory_value("webserver", "dag_orientation")
+ # this will only be set at serialization time
+ # it's only use is for determining the relative fileloc based only on the serialize dag
+ _processor_dags_folder: str | None = attrs.field(init=False, default=None)
+
# Override the default from parent class to use config
max_consecutive_failed_dag_runs: int = attrs.field()
@@ -435,17 +434,9 @@ class DAG(TaskSDKDag, LoggingMixin):
def _max_consecutive_failed_dag_runs_default(self):
return airflow_conf.getint("core", "max_consecutive_failed_dag_runs_per_dag")
- # this will only be set at serialization time
- # it's only use is for determining the relative fileloc based only on the serialize dag
- _processor_dags_folder: str | None = attrs.field(init=False, default=None)
-
- @has_on_success_callback.default
- def _has_on_success_callback(self) -> bool:
- return self.on_success_callback is not None
-
- @has_on_failure_callback.default
- def _has_on_failure_callback(self) -> bool:
- return self.on_failure_callback is not None
+ def validate(self):
+ super().validate()
+ self.validate_executor_field()
def validate_executor_field(self):
for task in self.tasks:
diff --git a/airflow/models/taskmixin.py b/airflow/models/taskmixin.py
index c3ab50e2e0b..fa76a3815cb 100644
--- a/airflow/models/taskmixin.py
+++ b/airflow/models/taskmixin.py
@@ -16,33 +16,13 @@
# under the License.
from __future__ import annotations
-from typing import TYPE_CHECKING, Any, Iterable
+from typing import TYPE_CHECKING
if TYPE_CHECKING:
- from airflow.models.operator import Operator
- from airflow.serialization.enums import DagAttributeTypes
from airflow.typing_compat import TypeAlias
import airflow.sdk.definitions.mixins
import airflow.sdk.definitions.node
DependencyMixin: TypeAlias = airflow.sdk.definitions.mixins.DependencyMixin
-
-
-class DAGNode(airflow.sdk.definitions.node.DAGNode):
- """
- A base class for a node in the graph of a workflow.
-
- A node may be an Operator or a Task Group, either mapped or unmapped.
- """
-
- def get_direct_relatives(self, upstream: bool = False) -> Iterable[Operator]:
- """Get list of the direct relatives to the current task, upstream or downstream."""
- if upstream:
- return self.upstream_list
- else:
- return self.downstream_list
-
- def serialize_for_task_group(self) -> tuple[DagAttributeTypes, Any]:
- """Serialize a task group's content; used by TaskGroupSerialization."""
- raise NotImplementedError()
+DAGNode: TypeAlias = airflow.sdk.definitions.node.DAGNode
diff --git a/task_sdk/src/airflow/sdk/definitions/baseoperator.py b/task_sdk/src/airflow/sdk/definitions/baseoperator.py
index d8e07c44b71..5796a6bbc1f 100644
--- a/task_sdk/src/airflow/sdk/definitions/baseoperator.py
+++ b/task_sdk/src/airflow/sdk/definitions/baseoperator.py
@@ -57,6 +57,7 @@ from airflow.task.priority_strategy import (
validate_and_load_priority_weight_strategy,
)
from airflow.utils import timezone
+from airflow.utils.setup_teardown import SetupTeardownContext
from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.types import AttributeRemoved
@@ -854,11 +855,11 @@ class BaseOperator(AbstractOperator, metaclass=BaseOperatorMeta):
)
self.template_fields = [self.template_fields]
- self._is_setup = False
- self._is_teardown = False
- # TODO: Task-SDK
- # if SetupTeardownContext.active:
- # SetupTeardownContext.update_context_map(self)
+ self.is_setup = False
+ self.is_teardown = False
+
+ if SetupTeardownContext.active:
+ SetupTeardownContext.update_context_map(self)
validate_instance_args(self, BASEOPERATOR_ARGS_EXPECTED_TYPES)
@@ -944,7 +945,7 @@ class BaseOperator(AbstractOperator, metaclass=BaseOperatorMeta):
result = cls.__new__(cls)
memo[id(self)] = result
- shallow_copy = cls.shallow_copy_attrs + cls._base_operator_shallow_copy_attrs
+ shallow_copy = tuple(cls.shallow_copy_attrs) + cls._base_operator_shallow_copy_attrs
for k, v in self.__dict__.items():
if k not in shallow_copy:
@@ -1173,8 +1174,6 @@ class BaseOperator(AbstractOperator, metaclass=BaseOperatorMeta):
"_BaseOperator__instantiated",
"_BaseOperator__init_kwargs",
"_BaseOperator__from_mapped",
- "_is_setup",
- "_is_teardown",
"_on_failure_fail_dagrun",
}
| { # Class level defaults need to be added to this list
diff --git a/task_sdk/src/airflow/sdk/definitions/dag.py b/task_sdk/src/airflow/sdk/definitions/dag.py
index e784639f8c9..111f51ce855 100644
--- a/task_sdk/src/airflow/sdk/definitions/dag.py
+++ b/task_sdk/src/airflow/sdk/definitions/dag.py
@@ -25,7 +25,7 @@ import os
import sys
import weakref
from collections import abc
-from collections.abc import Collection, Iterable, Iterator, MutableSet
+from collections.abc import Collection, Iterable, MutableSet
from datetime import datetime, timedelta
from inspect import signature
from re import Pattern
@@ -177,6 +177,20 @@ def _convert_access_control(value, self_: DAG):
return value
+def _convert_doc_md(doc_md: str | None) -> str | None:
+ if doc_md is None:
+ return doc_md
+
+ if doc_md.endswith(".md"):
+ try:
+ with open(doc_md) as fh:
+ return fh.read()
+ except FileNotFoundError:
+ return doc_md
+
+ return doc_md
+
+
def _all_after_dag_id_to_kw_only(cls, fields: list[attrs.Attribute]):
i = iter(fields)
f = next(i)
@@ -351,16 +365,17 @@ class DAG:
)
# sla_miss_callback: None | SLAMissCallback | list[SLAMissCallback] = None
catchup: bool = attrs.field(default=True, converter=bool)
- # on_success_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None
- # on_failure_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None
- doc_md: str | None = None
+ on_success_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None
+ on_failure_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None
+ doc_md: str | None = attrs.field(default=None, converter=_convert_doc_md)
params: ParamsDict = attrs.field(
# mypy doesn't really like passing the Converter object
default=None,
converter=attrs.Converter(_convert_params, takes_self=True), # type: ignore[misc, call-overload]
)
access_control: dict[str, dict[str, Collection[str]]] | dict[str, Collection[str]] | None = attrs.field(
- default=None, converter=attrs.Converter(_convert_access_control, takes_self=True)
+ default=None,
+ converter=attrs.Converter(_convert_access_control, takes_self=True), # type: ignore[misc, call-overload]
)
is_paused_upon_creation: bool | None = None
jinja_environment_kwargs: dict | None = None
@@ -368,7 +383,7 @@ class DAG:
tags: MutableSet[str] = attrs.field(factory=set, converter=_convert_tags)
owner_links: dict[str, str] = attrs.field(factory=dict)
auto_register: bool = attrs.field(default=True, converter=bool)
- fail_stop: bool = attrs.field(default=True, converter=bool)
+ fail_stop: bool = attrs.field(default=False, converter=bool)
dag_display_name: str = attrs.field(validator=attrs.validators.instance_of(str))
task_dict: dict[str, Operator] = attrs.field(factory=dict, init=False)
@@ -380,6 +395,9 @@ class DAG:
edge_info: dict[str, dict[str, EdgeInfoType]] = attrs.field(init=False, factory=dict)
+ has_on_success_callback: bool = attrs.field(init=False)
+ has_on_failure_callback: bool = attrs.field(init=False)
+
def __attrs_post_init__(self):
from airflow.utils import timezone
@@ -481,6 +499,23 @@ class DAG:
if tags and any(len(tag) > TAG_MAX_LEN for tag in tags):
raise ValueError(f"tag cannot be longer than {TAG_MAX_LEN} characters")
+ @max_active_runs.validator
+ def _validate_max_active_runs(self, _, max_active_runs):
+ if self.timetable.active_runs_limit is not None:
+ if self.timetable.active_runs_limit < self.max_active_runs:
+ raise ValueError(
+ f"Invalid max_active_runs: {type(self.timetable).__name__} "
+ f"requires max_active_runs <= {self.timetable.active_runs_limit}"
+ )
+
+ @has_on_success_callback.default
+ def _has_on_success_callback(self) -> bool:
+ return self.on_success_callback is not None
+
+ @has_on_failure_callback.default
+ def _has_on_failure_callback(self) -> bool:
+ return self.on_failure_callback is not None
+
def __repr__(self):
return f"<DAG: {self.dag_id}>"
@@ -522,18 +557,6 @@ class DAG:
_ = DagContext.pop()
- def get_doc_md(self, doc_md: str | None) -> str | None:
- if doc_md is None:
- return doc_md
-
- if doc_md.endswith(".md"):
- try:
- return open(doc_md).read()
- except FileNotFoundError:
- return doc_md
-
- return doc_md
-
def validate(self):
"""
Validate the DAG has a coherent setup.
@@ -543,6 +566,10 @@ class DAG:
self.timetable.validate()
self.validate_setup_teardown()
+ # We validate owner links on set, but since it's a dict it could be mutated without calling the
+ # setter. Validate again here
+ self._validate_owner_links(None, self.owner_links)
+
def validate_setup_teardown(self):
"""
Validate that setup and teardown tasks are configured properly.
@@ -966,20 +993,23 @@ class DAG:
"""
self.edge_info.setdefault(upstream_task_id, {})[downstream_task_id] = info
- def iter_invalid_owner_links(self) -> Iterator[tuple[str, str]]:
- """
- Parse a given link, and verifies if it's a valid URL, or a 'mailto' link.
+ @owner_links.validator
+ def _validate_owner_links(self, _, owner_links):
+ wrong_links = {}
- Returns an iterator of invalid (owner, link) pairs.
- """
- for owner, link in self.owner_links.items():
+ for owner, link in owner_links.items():
result = urlsplit(link)
if result.scheme == "mailto":
# netloc is not existing for 'mailto' link, so we are checking that the path is parsed
if not result.path:
- yield result.path, link
+ wrong_links[result.path] = link
elif not result.scheme or not result.netloc:
- yield owner, link
+ wrong_links[owner] = link
+ if wrong_links:
+ raise ValueError(
+ "Wrong link format was used for the owner. Use a valid link \n"
+ f"Bad formatted links are: {wrong_links}"
+ )
if TYPE_CHECKING:
@@ -1003,8 +1033,8 @@ if TYPE_CHECKING:
dagrun_timeout: timedelta | None = None,
# sla_miss_callback: Any = None,
catchup: bool = ...,
- # on_success_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None,
- # on_failure_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None,
+ on_success_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None,
+ on_failure_callback: None | DagStateChangeCallback | list[DagStateChangeCallback] = None,
doc_md: str | None = None,
params: ParamsDict | None = None,
access_control: dict[str, dict[str, Collection[str]]] | dict[str, Collection[str]] | None = None,
diff --git a/task_sdk/src/airflow/sdk/definitions/node.py b/task_sdk/src/airflow/sdk/definitions/node.py
index 7b877bbaf48..b3b519a07f0 100644
--- a/task_sdk/src/airflow/sdk/definitions/node.py
+++ b/task_sdk/src/airflow/sdk/definitions/node.py
@@ -211,6 +211,13 @@ class DAGNode(DependencyMixin, metaclass=ABCMeta):
else:
return self.downstream_task_ids
+ def get_direct_relatives(self, upstream: bool = False) -> Iterable[Operator]:
+ """Get list of the direct relatives to the current task, upstream or downstream."""
+ if upstream:
+ return self.upstream_list
+ else:
+ return self.downstream_list
+
def serialize_for_task_group(self) -> tuple[DagAttributeTypes, Any]:
"""Serialize a task group's content; used by TaskGroupSerialization."""
raise NotImplementedError()
diff --git a/task_sdk/tests/defintions/test_dag.py b/task_sdk/tests/defintions/test_dag.py
index c3b3bbdce4d..b2481a49b6a 100644
--- a/task_sdk/tests/defintions/test_dag.py
+++ b/task_sdk/tests/defintions/test_dag.py
@@ -248,6 +248,39 @@ class TestDag:
# Make sure we don't affect the original!
assert task.task_group.upstream_group_ids is not copied_task.task_group.upstream_group_ids
+ def test_dag_owner_links(self):
+ dag = DAG(
+ "dag",
+ schedule=None,
+ start_date=DEFAULT_DATE,
+ owner_links={"owner1": "https://mylink.com", "owner2": "mailto:someone@yoursite.com"},
+ )
+
+ assert dag.owner_links == {"owner1": "https://mylink.com", "owner2": "mailto:someone@yoursite.com"}
+
+ # Check wrong formatted owner link
+ with pytest.raises(ValueError, match="Wrong link format"):
+ DAG("dag", schedule=None, start_date=DEFAULT_DATE, owner_links={"owner1": "my-bad-link"})
+
+ dag = DAG("dag", schedule=None, start_date=DEFAULT_DATE)
+ dag.owner_links["owner1"] = "my-bad-link"
+ with pytest.raises(ValueError, match="Wrong link format"):
+ dag.validate()
+
+ def test_continuous_schedule_linmits_max_active_runs(self):
+ from airflow.timetables.simple import ContinuousTimetable
+
+ dag = DAG("continuous", start_date=DEFAULT_DATE, schedule="@continuous", max_active_runs=1)
+ assert isinstance(dag.timetable, ContinuousTimetable)
+ assert dag.max_active_runs == 1
+
+ dag = DAG("continuous", start_date=DEFAULT_DATE, schedule="@continuous", max_active_runs=0)
+ assert isinstance(dag.timetable, ContinuousTimetable)
+ assert dag.max_active_runs == 0
+
+ with pytest.raises(ValueError, match="ContinuousTimetable requires max_active_runs <= 1"):
+ dag = DAG("continuous", start_date=DEFAULT_DATE, schedule="@continuous", max_active_runs=25)
+
# Test some of the arg valiadtion. This is not all the validations we perform, just some of them.
@pytest.mark.parametrize(
@@ -255,6 +288,7 @@ class TestDag:
[
pytest.param("max_consecutive_failed_dag_runs", "not_an_int", id="max_consecutive_failed_dag_runs"),
pytest.param("dagrun_timeout", "not_an_int", id="dagrun_timeout"),
+ pytest.param("max_active_runs", "not_an_int", id="max_active_runs"),
],
)
def test_invalid_type_for_args(attr: str, value: Any):
diff --git a/tests/models/test_dag.py b/tests/models/test_dag.py
index c126e2d3039..4c1d8a67960 100644
--- a/tests/models/test_dag.py
+++ b/tests/models/test_dag.py
@@ -76,7 +76,6 @@ from airflow.templates import NativeEnvironment, SandboxedEnvironment
from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable
from airflow.timetables.simple import (
AssetTriggeredTimetable,
- ContinuousTimetable,
NullTimetable,
OnceTimetable,
)
@@ -1683,7 +1682,15 @@ class TestDag:
with dag:
check_task_2(check_task())
- dag.test()
+ dr = dag.test()
+
+ ti1 = dr.get_task_instance("check_task")
+ ti2 = dr.get_task_instance("check_task_2")
+
+ assert ti1
+ assert ti2
+ assert ti1.state == TaskInstanceState.FAILED
+ assert ti2.state == TaskInstanceState.UPSTREAM_FAILED
mock_handle_object_1.assert_called_with("task check_task failed...")
mock_handle_object_2.assert_called_with("dag test_local_testing_conn_file run failed...")
@@ -2121,22 +2128,6 @@ my_postgres_conn:
orm_dag_owners = session.query(DagOwnerAttributes).all()
assert not orm_dag_owners
- # Check wrong formatted owner link
- with pytest.raises(AirflowException):
- DAG("dag", schedule=None, start_date=DEFAULT_DATE, owner_links={"owner1": "my-bad-link"})
-
- def test_continuous_schedule_linmits_max_active_runs(self):
- dag = DAG("continuous", start_date=DEFAULT_DATE, schedule="@continuous", max_active_runs=1)
- assert isinstance(dag.timetable, ContinuousTimetable)
- assert dag.max_active_runs == 1
-
- dag = DAG("continuous", start_date=DEFAULT_DATE, schedule="@continuous", max_active_runs=0)
- assert isinstance(dag.timetable, ContinuousTimetable)
- assert dag.max_active_runs == 0
-
- with pytest.raises(AirflowException):
- dag = DAG("continuous", start_date=DEFAULT_DATE, schedule="@continuous", max_active_runs=25)
-
class TestDagModel:
def _clean(self):
@@ -2854,7 +2845,7 @@ def test_set_task_group_state(run_id, execution_date, session, dag_maker):
}
-def test_dag_teardowns_property_lists_all_teardown_tasks(dag_maker):
+def test_dag_teardowns_property_lists_all_teardown_tasks():
@setup
def setup_task():
return 1
@@ -2875,7 +2866,7 @@ def test_dag_teardowns_property_lists_all_teardown_tasks(dag_maker):
def mytask():
return 1
- with dag_maker() as dag:
+ with DAG("dag") as dag:
t1 = setup_task()
t2 = teardown_task()
t3 = teardown_task2()