This is an automated email from the ASF dual-hosted git repository. jason810496 pushed a commit to branch jason/lang-sdk-e2e/03e-lang-sdk-dag-bag in repository https://gitbox.apache.org/repos/asf/airflow.git
commit e0f332c645fc9234affbe12dde37c4f48b9d7f9c Author: ZHE YOU LIU <[email protected]> AuthorDate: Thu Oct 1 14:52:23 2026 +0000 Parse coordinator-claimed Dag files with their runtime in a Dag bag CoordinatorDagImporter.import_definition now runs the same LangSDKDagFileProcessorProcess as the Dag processor, for a Dag bag such as a CLI command's, and returns each Dag as a SerializedLangSDKDAG. A Dag bag has no API client, so the runtime's requests get an error there. The Dag bag skips the SDK-only checks and cluster policies for such a Dag, sync_bag_to_db does not store its Dags, and CLI commands that run a Dag refuse it. The recorded TypeScript and Java runtime payloads must pass validate_serialized_dag. --- .../src/airflow/cli/commands/dag_command.py | 9 +- airflow-core/src/airflow/dag_processing/dagbag.py | 79 ++- .../src/airflow/dag_processing/processor.py | 3 + .../src/airflow/serialization/definitions/dag.py | 8 +- airflow-core/src/airflow/utils/cli.py | 12 +- .../tests/unit/dag_processing/test_dagbag.py | 191 +++++- .../tests/unit/dag_processing/test_processor.py | 17 + airflow-core/tests/unit/utils/test_cli_util.py | 56 ++ .../src/airflow/sdk/coordinators/_dag_importer.py | 43 +- task-sdk/src/airflow/sdk/importers/base.py | 3 +- .../coordinators/fixtures/java_native.json | 151 +++++ .../coordinators/fixtures/ts_conformance.json | 661 +++++++++++++++++++++ .../task_sdk/coordinators/test_dag_importer.py | 93 ++- .../tests/task_sdk/importers/test_zip_importer.py | 3 +- 14 files changed, 1251 insertions(+), 78 deletions(-) diff --git a/airflow-core/src/airflow/cli/commands/dag_command.py b/airflow-core/src/airflow/cli/commands/dag_command.py index bd42dadf925..1f96179bfbb 100644 --- a/airflow-core/src/airflow/cli/commands/dag_command.py +++ b/airflow-core/src/airflow/cli/commands/dag_command.py @@ -69,6 +69,7 @@ if TYPE_CHECKING: from sqlalchemy.orm import Session from airflow import DAG + from airflow.dag_processing.dagbag import BaggedDAG from airflow.serialization.definitions.dag import SerializedDAG from airflow.timetables.base import DagRunInfo @@ -373,7 +374,7 @@ def _save_dot_to_file(dot: Dot, filename: str) -> None: print(f"File {filename} saved") -def _get_dagbag_dag_details(dag: DAG) -> dict: +def _get_dagbag_dag_details(dag: BaggedDAG) -> dict: """Return a dagbag dag details dict.""" from airflow.serialization.encoders import coerce_to_core_timetable @@ -599,7 +600,7 @@ def dag_list_dags(args, *, session: Session = NEW_SESSION) -> None: file=sys.stderr, ) - def get_dag_detail(dag: DAG) -> dict: + def get_dag_detail(dag: BaggedDAG) -> dict: if dag_model := DagModel.get_dagmodel(dag.dag_id, session=session): dag_detail = DAGResponse.model_validate(dag_model, from_attributes=True).model_dump() else: @@ -608,7 +609,9 @@ def dag_list_dags(args, *, session: Session = NEW_SESSION) -> None: return dag_detail return {col: dag_detail[col] for col in cols if col in DAG_DETAIL_FIELDS} - def filter_dags_by_bundle(dags: Iterable[DAG], bundle_names: list[str] | None) -> Iterable[DAG]: + def filter_dags_by_bundle( + dags: Iterable[BaggedDAG], bundle_names: list[str] | None + ) -> Iterable[BaggedDAG]: """Filter DAGs based on the specified bundle name, if provided.""" if not bundle_names: return dags diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 4e77a94e65d..222a11c793c 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -25,7 +25,7 @@ import warnings from collections.abc import Iterator from datetime import timedelta from pathlib import Path -from typing import TYPE_CHECKING, Any, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias from tabulate import tabulate @@ -46,6 +46,7 @@ from airflow.executors.executor_loader import ExecutorLoader from airflow.listeners.listener import get_listener_manager from airflow.models.pool import Pool from airflow.sdk.importers import DagImportError, get_importer_registry +from airflow.serialization.definitions.dag import SerializedLangSDKDAG from airflow.serialization.definitions.notset import NOTSET, ArgNotSet, is_arg_set from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.utils.file import correct_maybe_zipped, find_enclosing_file @@ -59,6 +60,9 @@ if TYPE_CHECKING: from airflow.models.dagwarning import DagWarning from airflow.sdk.importers import AbstractDagImporter, DagDefinition, DagImportWarning, DagSourceCode + BaggedDAG: TypeAlias = DAG | SerializedLangSDKDAG + """A Dag a Dag bag holds: a Python Dag, or a native Dag from a Lang-SDK runtime.""" + class FileLoadStat(NamedTuple): """ @@ -207,7 +211,7 @@ class DagBag(LoggingMixin): dag_folder = dag_folder or settings.DAGS_FOLDER self.dag_folder = dag_folder - self.dags: dict[str, DAG] = {} + self.dags: dict[str, BaggedDAG] = {} # The freshness token of each definition when we last imported it, keyed by its fileloc self.file_last_changed: dict[str, str] = {} # Store import errors with relative file paths as keys (relative to bundle_path) @@ -321,7 +325,7 @@ class DagBag(LoggingMixin): def _process_definition( self, importer: AbstractDagImporter, definition: DagDefinition, *, only_if_updated: bool - ) -> list[DAG]: + ) -> list[BaggedDAG]: """Import a Dag definition and bag the Dags it defines.""" fileloc = repr(definition) freshness_token = definition.freshness_token @@ -359,9 +363,11 @@ class DagBag(LoggingMixin): dag.fileloc = fileloc dag.relative_fileloc = self._get_relative_fileloc(fileloc) dag.bundle_name = self.bundle_name - dag.validate() - _validate_executor_fields(dag, self.bundle_name) - _assign_default_team_pools(dag, self.bundle_name) + # The Dag processor does not run these on a Lang-SDK Dag, so a Dag bag does not either. + if not isinstance(dag, SerializedLangSDKDAG): + dag.validate() + _validate_executor_fields(dag, self.bundle_name) + _assign_default_team_pools(dag, self.bundle_name) self.bag_dag(dag=dag) bagged_dags.append(dag) except AirflowClusterPolicySkipDag: @@ -432,36 +438,41 @@ class DagBag(LoggingMixin): return str(Path(filepath).relative_to(self.bundle_path)) return filepath - def bag_dag(self, dag: DAG): + def bag_dag(self, dag: BaggedDAG): """ Add the DAG into the bag. + A ``SerializedLangSDKDAG`` is only checked for a duplicate id. + ``validate_serialized_dag`` already rejected a cycle in it, and the Dag processor does not run + cluster policies on it either. + :raises: AirflowDagCycleException if a cycle is detected. :raises: AirflowDagDuplicatedIdException if this dag already exists in the bag. """ - dag.check_cycle() - dag.resolve_template_files() - dag.last_loaded = timezone.utcnow() - - try: - settings.dag_policy(dag) - - for task in dag.tasks: - if getattr(task, "end_from_trigger", False) and get_listener_manager().has_listeners: - raise AirflowException( - "Listeners are not supported with end_from_trigger=True for deferrable operators. " - "Task %s in DAG %s has end_from_trigger=True with listeners from plugins. " - "Set end_from_trigger=False to use listeners.", - task.task_id, - dag.dag_id, - ) + if not isinstance(dag, SerializedLangSDKDAG): + dag.check_cycle() + dag.resolve_template_files() + dag.last_loaded = timezone.utcnow() - settings.task_policy(task) - except (AirflowClusterPolicyViolation, AirflowClusterPolicySkipDag): - raise - except Exception as e: - self.log.exception(e) - raise AirflowClusterPolicyError(e) + try: + settings.dag_policy(dag) + + for task in dag.tasks: + if getattr(task, "end_from_trigger", False) and get_listener_manager().has_listeners: + raise AirflowException( + "Listeners are not supported with end_from_trigger=True for deferrable operators. " + "Task %s in DAG %s has end_from_trigger=True with listeners from plugins. " + "Set end_from_trigger=False to use listeners.", + task.task_id, + dag.dag_id, + ) + + settings.task_policy(task) + except (AirflowClusterPolicyViolation, AirflowClusterPolicySkipDag): + raise + except Exception as e: + self.log.exception(e) + raise AirflowClusterPolicyError(e) from airflow.sdk.exceptions import AirflowDagCycleException try: @@ -589,8 +600,8 @@ def sync_bag_to_db( """ Save attributes about list of DAG to the DB. - Files that a Lang-SDK runtime parses are left out, with their import errors: the Dag processor - stores those. + Files that a Lang-SDK runtime parses are left out, with their Dags and import errors: the Dag + processor stores those. """ from airflow.dag_processing.collection import update_dag_parsing_results_in_db @@ -618,7 +629,11 @@ def sync_bag_to_db( update_dag_parsing_results_in_db( bundle_name, bundle_version, - [LazyDeserializedDAG.from_dag(dag) for dag in dagbag.dags.values()], + [ + LazyDeserializedDAG.from_dag(dag) + for dag in dagbag.dags.values() + if not isinstance(dag, SerializedLangSDKDAG) + ], import_errors, None, # file parsing duration is not well defined when parsing multiple files / multiple DAGs. dagbag.dag_warnings, diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index aa8a1edfa85..45e2d4c6b40 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -74,6 +74,7 @@ from airflow.sdk.execution_time.comms import ( from airflow.sdk.execution_time.supervisor import WatchedSubprocess, register_request_method from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance, _send_error_email_notification from airflow.sdk.importers import DagSourceCode # noqa: TC001 +from airflow.serialization.definitions.dag import SerializedDAG from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from airflow.utils.dag_version_inflation_checker import check_dag_file_stability from airflow.utils.file import iter_airflow_imports @@ -313,6 +314,8 @@ def _get_dag_with_task( ) dag = dagbag.dags[dag_id] + if isinstance(dag, SerializedDAG): + raise ValueError(f"DAG '{dag_id}' is a native Lang-SDK Dag, whose callbacks do not run in Python.") if task_id is not None: try: diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py b/airflow-core/src/airflow/serialization/definitions/dag.py index e6979fc44f3..6367e9d93f7 100644 --- a/airflow-core/src/airflow/serialization/definitions/dag.py +++ b/airflow-core/src/airflow/serialization/definitions/dag.py @@ -24,7 +24,7 @@ import itertools import operator import re import weakref -from typing import TYPE_CHECKING, TypedDict, cast, overload +from typing import TYPE_CHECKING, TypeAlias, TypedDict, cast, overload import attrs import structlog @@ -150,6 +150,8 @@ class SerializedDAG: # Only on serialized dag. last_loaded: datetime.datetime = attrs.field(init=False) + fileloc: str = attrs.field(init=False) + relative_fileloc: str | None = attrs.field(init=False) # Determine the relative fileloc based only on the serialize dag. _processor_dags_folder: str = attrs.field(init=False) @@ -1419,6 +1421,10 @@ class SerializedDAG: return empty +SerializedLangSDKDAG: TypeAlias = SerializedDAG +"""A Dag that a Lang-SDK runtime parsed and serialized, as a Dag bag holds it.""" + + @provide_session def _create_orm_dagrun( *, diff --git a/airflow-core/src/airflow/utils/cli.py b/airflow-core/src/airflow/utils/cli.py index e3491bf91f1..6c8eceffc90 100644 --- a/airflow-core/src/airflow/utils/cli.py +++ b/airflow-core/src/airflow/utils/cli.py @@ -44,6 +44,7 @@ from airflow.utils.platform import getuser, is_terminal_support_colors T = TypeVar("T", bound=Callable) if TYPE_CHECKING: + from airflow.dag_processing.dagbag import BaggedDAG from airflow.sdk import DAG from airflow.serialization.definitions.dag import SerializedDAG @@ -291,6 +292,13 @@ def get_bagged_dag(bundle_names: list | None, dag_id: str, dagfile_path: str | N """ from airflow.dag_processing.dagbag import BundleDagBag, sync_bag_to_db from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager + from airflow.serialization.definitions.dag import SerializedLangSDKDAG + + def check_python_dag(dag: BaggedDAG) -> DAG: + # TODO: Support running a Lang-SDK Dag directly from the CLI. + if isinstance(dag, SerializedLangSDKDAG): + raise SystemExit(f"Dag {dag_id!r} is a native Lang-SDK Dag, which this command cannot run.") + return dag manager = DagBundlesManager() for bundle_name in bundle_names or (): @@ -302,7 +310,7 @@ def get_bagged_dag(bundle_names: list | None, dag_id: str, dagfile_path: str | N bundle_name=bundle.name, ) if dag := dagbag.dags.get(dag_id): - return dag + return check_python_dag(dag) manager.sync_bundles_to_db() for bundle in manager.get_all_dag_bundles(): @@ -315,7 +323,7 @@ def get_bagged_dag(bundle_names: list | None, dag_id: str, dagfile_path: str | N ) sync_bag_to_db(dagbag, bundle.name, bundle.version) if dag := dagbag.dags.get(dag_id): - return dag + return check_python_dag(dag) if dag: break raise AirflowException( diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index ecfaef5612d..7bfe751de00 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -42,6 +42,8 @@ from airflow.dag_processing.dagbag import ( _validate_executor_fields, sync_bag_to_db, ) +from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess +from airflow.dag_processing.processor import DagFileParsingResult from airflow.exceptions import UnknownExecutorException from airflow.executors.executor_loader import ExecutorLoader from airflow.models.dag import DagModel @@ -60,12 +62,20 @@ from airflow.sdk.importers import ( find_file_dag_definitions, get_file_suffix, ) +from airflow.serialization.definitions.baseoperator import SerializedBaseOperator +from airflow.serialization.definitions.dag import SerializedDAG +from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH from tests_common.test_utils import db from tests_common.test_utils.config import conf_vars from unit import cluster_policies -from unit.dag_processing.fake_lang_sdk import fake_coordinator, write_native_file +from unit.dag_processing.fake_lang_sdk import ( + FakeCoordinator, + fake_coordinator, + play_runtime, + write_native_file, +) from unit.models import TEST_DAGS_FOLDER pytestmark = pytest.mark.db_test @@ -1534,19 +1544,170 @@ class TestBundlePathSysPath: assert sys.path == syspath_before -def test_sync_bag_to_db_leaves_native_files_to_the_dag_processor(tmp_path, session, testing_dag_bundle): - db.clear_db_import_errors() - write_native_file(tmp_path / "dags.native") - session.add(ParseImportError(bundle_name="testing", filename="dags.native", stacktrace="stored")) - session.commit() +def _serialize_native_dag(dag_id: str, fileloc: str, relative_fileloc: str) -> dict: + with DAG(dag_id, schedule=None) as dag: + BaseOperator(task_id="extract") + data = DagSerialization.to_dict(dag) + data["dag"].update(fileloc=fileloc, relative_fileloc=relative_fileloc) + return data + + [email protected] +def _native_runtime(**dag_ids_by_file: list[str]): + """Make the runtime of a ``.native`` file return the Dag ids given for that file's stem.""" + + def run(*, path, dag_file_rel_path, **kwargs): + dag_ids = dag_ids_by_file[Path(path).stem] + return DagFileParsingResult( + fileloc=os.fspath(path), + serialized_dags=[ + LazyDeserializedDAG(data=_serialize_native_dag(dag_id, os.fspath(path), dag_file_rel_path)) + for dag_id in dag_ids + ], + ) + + with fake_coordinator(), patch.object(LangSDKDagFileProcessorProcess, "run", side_effect=run): + yield + + +class TestCoordinatorParsedFiles: + """A file claimed by a coordinator's Dag importer is parsed by its runtime and bagged as SerializedDAGs.""" + + @staticmethod + def _reply(*dag_ids, **result): + def reply(request, comms): + dags = [] + for dag_id in dag_ids: + with DAG(dag_id, schedule=None) as dag: + BaseOperator(task_id="extract") + data = DagSerialization.to_dict(dag) + data["dag"].update(fileloc=request.file, relative_fileloc="sub/dags.native") + dags.append(LazyDeserializedDAG(data=data)) + return DagFileParsingResult(fileloc=request.file, serialized_dags=dags, **result) + + return reply + + @mock.patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_native_file_is_bagged_as_serialized_dags(self, mock_parse_dag, tmp_path): + mock_parse_dag.side_effect = play_runtime(self._reply("native_a", "native_b")) + (tmp_path / "sub").mkdir() + native = write_native_file(tmp_path / "sub" / "dags.native") + + with fake_coordinator(): + dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing") + + assert dagbag.import_errors == {} + assert sorted(dagbag.dag_ids) == ["native_a", "native_b"] + dag = dagbag.dags["native_a"] + assert isinstance(dag, SerializedDAG) + assert isinstance(dag.task_dict["extract"], SerializedBaseOperator) + assert (dag.fileloc, dag.relative_fileloc) == (os.fspath(native), "sub/dags.native") + assert dag.bundle_name == "testing" + + @patch("airflow.settings.task_policy", cluster_policies.task_policy) + @patch("airflow.settings.dag_policy", cluster_policies.dag_policy) + @patch("airflow.dag_processing.dagbag._validate_executor_fields", autospec=True) + def test_native_dags_skip_the_python_dag_checks(self, mock_validate_executor_fields, tmp_path): + write_native_file(tmp_path / "dags.native") + + with _native_runtime(dags=["native_a"]): + dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing") + + assert dagbag.import_errors == {} + assert dagbag.dag_ids == ["native_a"] + mock_validate_executor_fields.assert_not_called() + + def test_native_dag_with_the_id_of_a_python_dag_is_an_import_error(self, tmp_path): + write_native_file(tmp_path / "dags.native") + (tmp_path / "a_python.py").write_text( + "from airflow.sdk import DAG\nwith DAG('shared_id', schedule=None): pass\n" + ) + + with _native_runtime(dags=["shared_id"]): + dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing") + + assert isinstance(dagbag.dags["shared_id"], DAG) + assert dagbag.import_errors["dags.native"].startswith("AirflowDagDuplicatedIdException: ") - with fake_coordinator(): - dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing") - sync_bag_to_db(dagbag, "testing", None, session=session) + def test_sync_bag_to_db_leaves_native_files_to_the_dag_processor( + self, tmp_path, session, testing_dag_bundle + ): + db.clear_db_dags() + db.clear_db_serialized_dags() + db.clear_db_import_errors() + write_native_file(tmp_path / "dags.native") + write_native_file(tmp_path / "broken.native") + (tmp_path / "a_python.py").write_text( + "from airflow.sdk import DAG\nwith DAG('python_dag', schedule=None): pass\n" + ) + session.add(ParseImportError(bundle_name="testing", filename="dags.native", stacktrace="stored")) + session.commit() + + with _native_runtime(dags=["native_a"], broken=[]): + dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing") + dagbag.import_errors["broken.native"] = "cannot start the runtime here" + sync_bag_to_db(dagbag, "testing", None, session=session) + + assert isinstance(dagbag.dags["native_a"], SerializedDAG) + assert set(session.scalars(select(DagModel.dag_id))) == {"python_dag"} + assert {(e.filename, e.stacktrace) for e in session.scalars(select(ParseImportError))} == { + ("dags.native", "stored") + } + + def test_sync_bag_to_db_keeps_a_native_dag_whose_runtime_cannot_start( + self, tmp_path, session, testing_dag_bundle + ): + db.clear_db_dags() + native = write_native_file(tmp_path / "dags.native", command_error="no runtime") + session.add( + DagModel( + dag_id="native_dag", + bundle_name="testing", + fileloc=os.fspath(native), + relative_fileloc="dags.native", + is_stale=False, + ) + ) + session.commit() + + with fake_coordinator(): + dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing") + sync_bag_to_db(dagbag, "testing", None, session=session) + + assert dagbag.import_errors == { + "dags.native": "Cannot start the Lang-SDK runtime: FileNotFoundError: no runtime" + } + assert session.scalar(select(DagModel.is_stale).where(DagModel.dag_id == "native_dag")) is False + + @pytest.mark.parametrize( + ("spec", "reply", "message"), + [ + pytest.param( + {}, + lambda request, comms: None, + "The Lang-SDK runtime exited with code 0 without a parse result", + id="no-result", + ), + pytest.param( + {"schema_version": "1999-01-01"}, + None, + "Cannot start the Lang-SDK runtime: " + "ValueError: Version '1999-01-01' not found in supervisor schema bundle", + id="cannot-start", + ), + pytest.param({}, _reply(import_errors={"main.ts": "boom"}), "main.ts: boom", id="runtime-error"), + ], + ) + def test_runtime_failures_are_import_errors(self, tmp_path, spec, reply, message): + native = write_native_file(tmp_path / "dags.native", **spec) + + with ( + fake_coordinator(), + mock.patch.object(FakeCoordinator, "parse_dag", autospec=True, side_effect=play_runtime(reply)) + if reply + else contextlib.nullcontext(), + ): + dagbag = DagBag(dag_folder=os.fspath(native), bundle_path=tmp_path, bundle_name="testing") - assert dagbag.import_errors == { - "dags.native": "A native Lang-SDK Dag is parsed only by the Dag processor" - } - assert {(e.filename, e.stacktrace) for e in session.scalars(select(ParseImportError))} == { - ("dags.native", "stored") - } + assert dagbag.dags == {} + assert dagbag.import_errors == {"dags.native": message} diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 0cb61aa0286..06e988213dc 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -89,6 +89,7 @@ from airflow.sdk.execution_time.comms import ( ) from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance from airflow.sdk.importers import DagSourceCode +from airflow.serialization.serialized_objects import DagSerialization from airflow.utils.session import create_session from airflow.utils.state import TaskInstanceState @@ -1289,6 +1290,22 @@ class TestExecuteDagCallbacks: with pytest.raises(ValueError, match="DAG 'missing_dag' not found in DagBag"): _execute_dag_callbacks(dagbag, request, log) + def test_execute_dag_callbacks_native_dag(self, tmp_path): + dagbag = DagBag(dag_folder=tmp_path) + native_dag = DAG("native_dag", schedule=None) + dagbag.dags["native_dag"] = DagSerialization.from_dict(DagSerialization.to_dict(native_dag)) + request = DagCallbackRequest( + filepath="dags.native", + dag_id="native_dag", + run_id="test_run", + bundle_name="testing", + bundle_version=None, + is_failure_callback=True, + ) + + with pytest.raises(ValueError, match="DAG 'native_dag' is a native Lang-SDK Dag"): + _execute_dag_callbacks(dagbag, request, structlog.get_logger()) + @pytest.mark.parametrize( ("xcom_operation", "expected_message_type", "expected_message", "mock_response"), [ diff --git a/airflow-core/tests/unit/utils/test_cli_util.py b/airflow-core/tests/unit/utils/test_cli_util.py index 3539f8b3e52..159bccb436f 100644 --- a/airflow-core/tests/unit/utils/test_cli_util.py +++ b/airflow-core/tests/unit/utils/test_cli_util.py @@ -32,15 +32,22 @@ from sqlalchemy import select import airflow from airflow import settings from airflow._shared.timezones import timezone +from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess +from airflow.dag_processing.processor import DagFileParsingResult from airflow.exceptions import AirflowException from airflow.models.dag import DagModel from airflow.models.dagbundle import DagBundleModel from airflow.models.log import Log from airflow.models.team import Team +from airflow.sdk import DAG, BaseOperator +from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG from airflow.utils import cli, cli_action_loggers from airflow.utils.cli import _search_for_dag_file +from airflow.utils.session import create_session +from tests_common.test_utils import db from tests_common.test_utils.config import conf_vars +from unit.dag_processing.fake_lang_sdk import fake_coordinator, write_native_file # Mark entire module as db_test because ``action_cli`` wrapper still could use DB on callbacks: # - ``cli_action_loggers.on_pre_execution`` @@ -342,6 +349,55 @@ def test__search_for_dags_file(): assert _search_for_dag_file("any/hi/__init__.py") is None [email protected]("bundle_names", [["testing"], None], ids=["named-bundle", "every-bundle"]) [email protected]("airflow.dag_processing.dagbag.sync_bag_to_db", autospec=True) [email protected]("airflow.dag_processing.dagbag.BundleDagBag", autospec=True) [email protected](cli, "DagBundlesManager", autospec=True) +def test_get_bagged_dag_refuses_a_native_dag(mock_manager, mock_bag, mock_sync, bundle_names): + mock_manager.return_value.get_all_dag_bundles.return_value = [mock.MagicMock()] + native = DagSerialization.from_dict(DagSerialization.to_dict(DAG("native", schedule=None))) + mock_bag.return_value.dags = {"native": native} + + with pytest.raises(SystemExit, match="is a native Lang-SDK Dag"): + cli.get_bagged_dag(bundle_names, "native") + + [email protected] +def _clear_db_dags(): + db.clear_db_dags() + db.clear_db_serialized_dags() + yield + db.clear_db_dags() + db.clear_db_serialized_dags() + + [email protected]("_clear_db_dags", "testing_dag_bundle") [email protected](LangSDKDagFileProcessorProcess, "run", autospec=True) [email protected](cli, "DagBundlesManager", autospec=True) +def test_get_bagged_dag_returns_a_python_dag_bagged_with_a_native_one(mock_manager, mock_run, tmp_path): + (tmp_path / "python_dag.py").write_text( + "from airflow.sdk import DAG\nwith DAG('python_dag', schedule=None): pass\n" + ) + native_file = write_native_file(tmp_path / "dags.native") + with DAG("native_dag", schedule=None) as native_dag: + BaseOperator(task_id="extract") + mock_run.return_value = DagFileParsingResult( + fileloc=os.fspath(native_file), serialized_dags=[LazyDeserializedDAG.from_dag(native_dag)] + ) + bundle = mock.MagicMock(path=tmp_path, version=None) + bundle.name = "testing" + mock_manager.return_value.get_all_dag_bundles.return_value = [bundle] + + with fake_coordinator(): + dag = cli.get_bagged_dag(None, "python_dag") + + assert isinstance(dag, DAG) + assert dag.dag_id == "python_dag" + mock_run.assert_called_once() + with create_session() as session: + assert set(session.scalars(select(DagModel.dag_id))) == {"python_dag"} + + def test_validate_dag_bundle_arg(): with pytest.raises(SystemExit, match="Bundles not found: (x, y)|(y, x)"): cli.validate_dag_bundle_arg(["x", "y", "dags-folder"]) diff --git a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py index 048c209b5e0..b7efebf7363 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py @@ -21,6 +21,8 @@ from __future__ import annotations from typing import TYPE_CHECKING, ClassVar +import structlog + from airflow.sdk.importers.base import ( AbstractDagImporter, DagImportError, @@ -33,18 +35,23 @@ if TYPE_CHECKING: from collections.abc import Iterator from pathlib import Path + from structlog.typing import FilteringBoundLogger + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.coordinators._subprocess import SubprocessCoordinator from airflow.sdk.importers.base import DagDefinition +log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.dag_importer") + class CoordinatorDagImporter(AbstractDagImporter[FilesystemDagDefinition]): """ - Claim the native Dag files of a coordinator's artifacts, which the coordinator's runtime parses. + Import the native Dags of a coordinator's artifacts by running the coordinator's runtime. The Dag processor does not call :meth:`import_definition`: it runs the runtime itself and stores - the Dags the runtime serialized. A Dag bag, such as a CLI command's, reports such a file as an - import error. + the Dags the runtime serialized. This method serves a Dag bag, such as a CLI command's, and returns + each Dag as a ``SerializedLangSDKDAG``, the ``SerializedDAG`` the scheduler loads. Its tasks run only through the coordinator, + never in Python. Subclasses set :attr:`artifact_suffix` and :attr:`supported_extensions`, and implement :meth:`get_source_code`. @@ -78,13 +85,27 @@ class CoordinatorDagImporter(AbstractDagImporter[FilesystemDagDefinition]): def import_definition( self, definition: FilesystemDagDefinition, bundle: BaseDagBundle ) -> DagImportResult: - """Report that only the Dag processor parses *definition*.""" - return DagImportResult( - definition=definition, - errors=[ + from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess # noqa: SDK002 + from airflow.serialization.serialized_objects import DagSerialization # noqa: SDK002 + + source_reference = repr(definition) + bundle_path = bundle.path or definition.path.parent + relative_loc = definition.get_relative_loc(bundle_path) + result = DagImportResult(definition=definition) + parsing_result = LangSDKDagFileProcessorProcess.run( + path=definition.path, + bundle_path=bundle_path, + bundle_name=bundle.name, + dag_file_rel_path=relative_loc, + logger=log, + ) + for key, message in (parsing_result.import_errors or {}).items(): + result.errors.append( DagImportError( - source_reference=repr(definition), - message="A native Lang-SDK Dag is parsed only by the Dag processor", + source_reference=source_reference, + message=message if key == relative_loc else f"{key}: {message}", ) - ], - ) + ) + # The runtime process validated each Dag, and moved one that fails into import_errors. + result.dags.extend(DagSerialization.from_dict(dag.data) for dag in parsing_result.serialized_dags) + return result diff --git a/task-sdk/src/airflow/sdk/importers/base.py b/task-sdk/src/airflow/sdk/importers/base.py index 1bb6b16414f..c92ac93f44d 100644 --- a/task-sdk/src/airflow/sdk/importers/base.py +++ b/task-sdk/src/airflow/sdk/importers/base.py @@ -38,6 +38,7 @@ if TYPE_CHECKING: from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk import DAG + from airflow.serialization.definitions.dag import SerializedLangSDKDAG # noqa: SDK002 log = logging.getLogger(__name__) @@ -187,7 +188,7 @@ class DagImportResult: """Result of importing DAGs from a definition.""" definition: DagDefinition | None = None - dags: list[DAG] = field(default_factory=list) + dags: list[DAG | SerializedLangSDKDAG] = field(default_factory=list) errors: list[DagImportError] = field(default_factory=list) skipped_definitions: list[DagDefinition] = field(default_factory=list) warnings: list[DagImportWarning] = field(default_factory=list) diff --git a/task-sdk/tests/task_sdk/coordinators/fixtures/java_native.json b/task-sdk/tests/task_sdk/coordinators/fixtures/java_native.json new file mode 100644 index 00000000000..50314784b10 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/fixtures/java_native.json @@ -0,0 +1,151 @@ +{ + "type": "DagFileParsingResult", + "fileloc": "/bundles/app/dags.jar", + "serialized_dags": [ + { + "data": { + "__version": 3, + "dag": { + "dag_id": "java_native", + "fileloc": "/bundles/app/dags.jar", + "relative_fileloc": "dags.jar", + "timezone": "UTC", + "timetable": { + "__type": "airflow.timetables.trigger.CronTriggerTimetable", + "__var": { + "expression": "@daily", + "timezone": "UTC", + "interval": 0.0, + "run_immediately": false + } + }, + "tasks": [ + { + "__type": "operator", + "__var": { + "task_id": "extract", + "task_type": "DumpExtract", + "_task_module": "org.apache.airflow.sdk.execution", + "language": "java", + "template_fields": [], + "retries": 2, + "queue": "java", + "retry_delay": 600.0, + "execution_timeout": 90.0, + "downstream_task_ids": [ + "load" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "load", + "task_type": "DumpLoad", + "_task_module": "org.apache.airflow.sdk.execution", + "language": "java", + "template_fields": [], + "queue": "java", + "trigger_rule": "all_done" + } + } + ], + "dag_dependencies": [], + "task_group": { + "_group_id": null, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "extract": [ + "operator", + "extract" + ], + "load": [ + "operator", + "load" + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + }, + "edge_info": {}, + "params": [], + "deadline": null, + "allowed_run_types": null, + "description": "java native", + "start_date": 1767225600.0, + "dagrun_timeout": 300.0, + "tags": [ + "a", + "b" + ], + "max_active_tasks": 16, + "max_active_runs": 3, + "max_consecutive_failed_dag_runs": 0, + "catchup": true, + "disable_bundle_versioning": false + } + } + }, + { + "data": { + "__version": 3, + "dag": { + "dag_id": "java_once", + "fileloc": "/bundles/app/dags.jar", + "relative_fileloc": "dags.jar", + "timezone": "UTC", + "timetable": { + "__type": "airflow.timetables.simple.OnceTimetable", + "__var": {} + }, + "tasks": [ + { + "__type": "operator", + "__var": { + "task_id": "solo", + "task_type": "DumpExtract", + "_task_module": "org.apache.airflow.sdk.execution", + "language": "java", + "template_fields": [] + } + } + ], + "dag_dependencies": [], + "task_group": { + "_group_id": null, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "solo": [ + "operator", + "solo" + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + }, + "edge_info": {}, + "params": [], + "deadline": null, + "allowed_run_types": null, + "max_active_tasks": 16, + "max_active_runs": 16, + "max_consecutive_failed_dag_runs": 0, + "catchup": false, + "disable_bundle_versioning": false + } + } + } + ] +} diff --git a/task-sdk/tests/task_sdk/coordinators/fixtures/ts_conformance.json b/task-sdk/tests/task_sdk/coordinators/fixtures/ts_conformance.json new file mode 100644 index 00000000000..d2af4fd8732 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/fixtures/ts_conformance.json @@ -0,0 +1,661 @@ +{ + "conformance_minimal": { + "__version": 3, + "dag": { + "dag_id": "conformance_minimal", + "fileloc": "/bundles/app/bundle.mjs", + "relative_fileloc": "bundle.mjs", + "timezone": "UTC", + "timetable": { + "__type": "airflow.timetables.simple.NullTimetable", + "__var": {} + }, + "tasks": [ + { + "__type": "operator", + "__var": { + "task_id": "solo", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true + } + } + ], + "dag_dependencies": [], + "task_group": { + "_group_id": null, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "solo": [ + "operator", + "solo" + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + }, + "edge_info": {}, + "params": [], + "deadline": null, + "allowed_run_types": null + } + }, + "conformance_cron": { + "__version": 3, + "dag": { + "dag_id": "conformance_cron", + "fileloc": "/bundles/app/bundle.mjs", + "relative_fileloc": "bundle.mjs", + "timezone": "UTC", + "timetable": { + "__type": "airflow.timetables.trigger.CronTriggerTimetable", + "__var": { + "expression": "0 3 * * *", + "timezone": "UTC", + "interval": 0, + "run_immediately": false + } + }, + "tasks": [ + { + "__type": "operator", + "__var": { + "task_id": "extract", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "owner": "data-team", + "start_date": 1769904000, + "end_date": 1795996800, + "trigger_rule": "all_done", + "depends_on_past": true, + "ignore_first_depends_on_past": true, + "wait_for_past_depends_before_skipping": true, + "wait_for_downstream": true, + "retries": 2, + "queue": "typescript", + "pool": "tiny", + "pool_slots": 2, + "execution_timeout": 120, + "retry_delay": 600, + "retry_exponential_backoff": 2, + "max_retry_delay": 900, + "priority_weight": 5, + "weight_rule": "upstream", + "executor": "LocalExecutor", + "do_xcom_push": false, + "doc_md": "extracts things", + "map_index_template": "{{ task.task_id }}", + "max_active_tis_per_dag": 3, + "max_active_tis_per_dagrun": 4, + "downstream_task_ids": [ + "transform" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "transform", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "_arg_bindings": [ + { + "name": "from_extract", + "kind": "xcom", + "task_id": "extract" + } + ] + } + } + ], + "dag_dependencies": [], + "task_group": { + "_group_id": null, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "extract": [ + "operator", + "extract" + ], + "transform": [ + "operator", + "transform" + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + }, + "edge_info": {}, + "params": [], + "deadline": null, + "allowed_run_types": null, + "description": "a Dag that sets everything", + "start_date": 1767225600, + "end_date": 1798759815, + "tags": [ + "alpha", + "gamma" + ], + "dag_display_name": "Conformance Cron", + "doc_md": "# notes", + "max_active_tasks": 8, + "max_active_runs": 3, + "max_consecutive_failed_dag_runs": 2, + "dagrun_timeout": 300, + "catchup": true, + "render_template_as_native_obj": true, + "disable_bundle_versioning": true, + "is_paused_upon_creation": true + } + }, + "conformance_once": { + "__version": 3, + "dag": { + "dag_id": "conformance_once", + "fileloc": "/bundles/app/bundle.mjs", + "relative_fileloc": "bundle.mjs", + "timezone": "UTC", + "timetable": { + "__type": "airflow.timetables.simple.OnceTimetable", + "__var": {} + }, + "tasks": [ + { + "__type": "operator", + "__var": { + "task_id": "seed", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "downstream_task_ids": [ + "alpha", + "beta" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "beta", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "_arg_bindings": [ + { + "name": "from_seed", + "kind": "xcom", + "task_id": "seed" + } + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "alpha", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "_arg_bindings": [ + { + "name": "from_seed", + "kind": "xcom", + "task_id": "seed" + } + ] + } + } + ], + "dag_dependencies": [], + "task_group": { + "_group_id": null, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "seed": [ + "operator", + "seed" + ], + "beta": [ + "operator", + "beta" + ], + "alpha": [ + "operator", + "alpha" + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + }, + "edge_info": {}, + "params": [], + "deadline": null, + "allowed_run_types": null, + "fail_fast": true + } + }, + "conformance_continuous": { + "__version": 3, + "dag": { + "dag_id": "conformance_continuous", + "fileloc": "/bundles/app/bundle.mjs", + "relative_fileloc": "bundle.mjs", + "timezone": "UTC", + "timetable": { + "__type": "airflow.timetables.simple.ContinuousTimetable", + "__var": {} + }, + "tasks": [ + { + "__type": "operator", + "__var": { + "task_id": "watch", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true + } + } + ], + "dag_dependencies": [], + "task_group": { + "_group_id": null, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "watch": [ + "operator", + "watch" + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + }, + "edge_info": {}, + "params": [], + "deadline": null, + "allowed_run_types": null, + "max_active_runs": 1 + } + }, + "conformance_diamond": { + "__version": 3, + "dag": { + "dag_id": "conformance_diamond", + "fileloc": "/bundles/app/bundle.mjs", + "relative_fileloc": "bundle.mjs", + "timezone": "UTC", + "timetable": { + "__type": "airflow.timetables.trigger.CronTriggerTimetable", + "__var": { + "expression": "0 0 * * *", + "timezone": "UTC", + "interval": 0, + "run_immediately": false + } + }, + "tasks": [ + { + "__type": "operator", + "__var": { + "task_id": "root", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "downstream_task_ids": [ + "left", + "right" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "right", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "_arg_bindings": [ + { + "name": "from_root", + "kind": "xcom", + "task_id": "root" + } + ], + "downstream_task_ids": [ + "join" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "left", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "_arg_bindings": [ + { + "name": "from_root", + "kind": "xcom", + "task_id": "root" + } + ], + "downstream_task_ids": [ + "join" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "join", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "_arg_bindings": [ + { + "name": "from_right", + "kind": "xcom", + "task_id": "right" + }, + { + "name": "from_left", + "kind": "xcom", + "task_id": "left" + } + ] + } + } + ], + "dag_dependencies": [], + "task_group": { + "_group_id": null, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "root": [ + "operator", + "root" + ], + "right": [ + "operator", + "right" + ], + "left": [ + "operator", + "left" + ], + "join": [ + "operator", + "join" + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + }, + "edge_info": {}, + "params": [], + "deadline": null, + "allowed_run_types": null + } + }, + "conformance_groups": { + "__version": 3, + "dag": { + "dag_id": "conformance_groups", + "fileloc": "/bundles/app/bundle.mjs", + "relative_fileloc": "bundle.mjs", + "timezone": "UTC", + "timetable": { + "__type": "airflow.timetables.simple.OnceTimetable", + "__var": {} + }, + "tasks": [ + { + "__type": "operator", + "__var": { + "task_id": "extract", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "downstream_task_ids": [ + "staging.stage" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "staging.stage", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "downstream_task_ids": [ + "staging.checks.nulls" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "staging.checks.nulls", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "_arg_bindings": [ + { + "name": "from_staging.stage", + "kind": "xcom", + "task_id": "staging.stage" + } + ], + "downstream_task_ids": [ + "publish.push" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "publish.push", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "downstream_task_ids": [ + "load" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "load", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true, + "downstream_task_ids": [ + "notify" + ] + } + }, + { + "__type": "operator", + "__var": { + "task_id": "notify", + "task_type": "TypeScriptOperator", + "_task_module": "airflow.sdk.coordinators.node", + "language": "typescript", + "template_fields": [], + "is_stub": true + } + } + ], + "dag_dependencies": [], + "task_group": { + "_group_id": null, + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "extract": [ + "operator", + "extract" + ], + "load": [ + "operator", + "load" + ], + "notify": [ + "operator", + "notify" + ], + "staging": [ + "taskgroup", + { + "_group_id": "staging", + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "staging.stage": [ + "operator", + "staging.stage" + ], + "staging.checks": [ + "taskgroup", + { + "_group_id": "checks", + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "staging.checks.nulls": [ + "operator", + "staging.checks.nulls" + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + } + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [ + "publish" + ], + "upstream_task_ids": [ + "extract" + ], + "downstream_task_ids": [] + } + ], + "publish": [ + "taskgroup", + { + "_group_id": "publish", + "group_display_name": "", + "prefix_group_id": true, + "tooltip": "", + "ui_color": "CornflowerBlue", + "ui_fgcolor": "#000", + "children": { + "publish.push": [ + "operator", + "publish.push" + ] + }, + "upstream_group_ids": [ + "staging" + ], + "downstream_group_ids": [], + "upstream_task_ids": [ + "staging.checks.nulls" + ], + "downstream_task_ids": [ + "load" + ] + } + ] + }, + "upstream_group_ids": [], + "downstream_group_ids": [], + "upstream_task_ids": [], + "downstream_task_ids": [] + }, + "edge_info": {}, + "params": [], + "deadline": null, + "allowed_run_types": null + } + } +} diff --git a/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py b/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py index abbada9cba7..6a0a9c970ee 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py +++ b/task-sdk/tests/task_sdk/coordinators/test_dag_importer.py @@ -17,14 +17,23 @@ # under the License. from __future__ import annotations +import copy +import json +from pathlib import Path from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import ANY, MagicMock, patch import pytest +from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess +from airflow.dag_processing.processor import DagFileParsingResult from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter from airflow.sdk.coordinators._subprocess import SubprocessCoordinator from airflow.sdk.importers import DagSourceCode, FilesystemDagDefinition +from airflow.serialization.definitions.dag import SerializedDAG +from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG + +FIXTURES = Path(__file__).parent / "fixtures" class _BundleImporter(CoordinatorDagImporter): @@ -38,6 +47,33 @@ class _BundleImporter(CoordinatorDagImporter): return not definition.path.name.startswith("skip") +def _read_payloads(path: Path) -> list[dict]: + recorded = json.loads(path.read_text()) + if "serialized_dags" in recorded: + return [serialized["data"] for serialized in recorded["serialized_dags"]] + # What the TS SDK's tests/conformance/serialize_typescript.ts writes: payloads keyed by Dag id. + return list(recorded.values()) + + +def _load_payloads() -> list: + return [ + pytest.param(data, id=f"{path.stem}-{data['dag']['dag_id']}") + for path in sorted(FIXTURES.glob("*.json")) + for data in _read_payloads(path) + ] + + +def _get_payload(dag_id: str) -> dict: + return copy.deepcopy( + next(p.values[0] for p in _load_payloads() if p.values[0]["dag"]["dag_id"] == dag_id) + ) + + [email protected]("data", _load_payloads()) +def test_recorded_runtime_payloads_are_valid(data): + DagSerialization.validate_serialized_dag(data) + + @pytest.fixture def importer() -> _BundleImporter: return _BundleImporter(coordinator=MagicMock(spec=SubprocessCoordinator)) @@ -57,14 +93,47 @@ def test_lists_only_its_artifacts(importer, tmp_path): assert [d.path.name for d in definitions] == ["main.min.mjs"] -def test_import_definition_reports_that_only_the_dag_processor_parses_it(importer, tmp_path): - bundle_file = tmp_path / "main.min.mjs" - bundle_file.write_text("") - definition = FilesystemDagDefinition(bundle_file) - - result = importer.import_definition(definition, SimpleNamespace(name="testing", path=tmp_path)) - - assert result.dags == [] - assert [error.message for error in result.errors] == [ - "A native Lang-SDK Dag is parsed only by the Dag processor" - ] +class TestImportDefinition: + @staticmethod + def _import(importer, tmp_path): + definition = FilesystemDagDefinition(tmp_path / "main.min.mjs") + return importer.import_definition(definition, SimpleNamespace(name="testing", path=tmp_path)) + + @patch.object(LangSDKDagFileProcessorProcess, "run", autospec=True) + def test_returns_the_dags_the_runtime_serialized(self, mock_run, importer, tmp_path): + mock_run.return_value = DagFileParsingResult( + fileloc=str(tmp_path / "main.min.mjs"), + serialized_dags=[LazyDeserializedDAG(data=_get_payload("conformance_minimal"))], + import_errors={"main.min.mjs": "one failed", "main.ts": "two failed"}, + ) + + result = self._import(importer, tmp_path) + + [dag] = result.dags + assert isinstance(dag, SerializedDAG) + assert (dag.dag_id, dag.task_ids) == ("conformance_minimal", ["solo"]) + assert [(e.source_reference, e.message) for e in result.errors] == [ + (str(tmp_path / "main.min.mjs"), "one failed"), + (str(tmp_path / "main.min.mjs"), "main.ts: two failed"), + ] + mock_run.assert_called_once_with( + path=tmp_path / "main.min.mjs", + bundle_path=tmp_path, + bundle_name="testing", + dag_file_rel_path="main.min.mjs", + logger=ANY, + ) + + @pytest.mark.parametrize("data", _load_payloads()) + @patch.object(LangSDKDagFileProcessorProcess, "run", autospec=True) + def test_imports_every_recorded_runtime_payload(self, mock_run, importer, tmp_path, data): + mock_run.return_value = DagFileParsingResult( + fileloc=data["dag"]["fileloc"], serialized_dags=[LazyDeserializedDAG(data=data)] + ) + + [dag] = self._import(importer, tmp_path).dags + + assert {t.task_id: t.downstream_task_ids for t in dag.tasks} == { + t["__var"]["task_id"]: set(t["__var"].get("downstream_task_ids", [])) + for t in data["dag"]["tasks"] + } diff --git a/task-sdk/tests/task_sdk/importers/test_zip_importer.py b/task-sdk/tests/task_sdk/importers/test_zip_importer.py index 3d4a52bc696..325f8578c4b 100644 --- a/task-sdk/tests/task_sdk/importers/test_zip_importer.py +++ b/task-sdk/tests/task_sdk/importers/test_zip_importer.py @@ -43,6 +43,7 @@ from airflow.sdk.importers.python_importer import PythonDagImporter if TYPE_CHECKING: from airflow.sdk import DAG + from airflow.serialization.definitions.dag import SerializedDAG class CustomInternalNonExtensionImporter(AbstractDagImporter): @@ -67,7 +68,7 @@ class CustomInternalNonExtensionImporter(AbstractDagImporter): def _import_all( importer: AbstractDagImporter[DagDefinition], bundle, -) -> tuple[list[DAG], list[DagImportError]]: +) -> tuple[list[DAG | SerializedDAG], list[DagImportError]]: """Enumerate an importer's definitions and import each, aggregating dags/errors.""" dags, errors = [], [] for item in importer.list_dag_definitions(bundle):
