Parse native Lang-SDK Dags with their coordinator's runtime

airflow-core/adr/lang-sdk/0004-dag-parsing.md

@@ -216,9 +216,11 @@ The language runtime must produce a `DagFileParsingResult` that matches Python A
 | `start_date` | float (epoch) | if set | Unwrapped from `__type`/`__var` |
 | `end_date` | float (epoch) | if set | Unwrapped from `__type`/`__var` |
 | `tags` | list | if non-empty | Unwrapped from `__type`/`__var` |
-| `catchup` | bool | if `true` | |
-| `max_active_tasks` | int | if non-default | |
-| `max_active_runs` | int | if non-default | |
+| `catchup` | bool | if set | Airflow fills an unset field from its config |
+| `max_active_tasks` | int | if set | Airflow fills an unset field from its config |
+| `max_active_runs` | int | if set | Airflow fills an unset field from its config |
+| `max_consecutive_failed_dag_runs` | int | if set | Airflow fills an unset field from its config |
+| `disable_bundle_versioning` | bool | if set | Airflow fills an unset field from its config |
 
 **Timetable encoding:**
 

airflow-core/src/airflow/cli/commands/dag_command.py

@@ -69,6 +69,7 @@
     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 get_dag_detail(dag: DAG) -> dict:
             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

airflow-core/src/airflow/dag_processing/dagbag.py

@@ -25,14 +25,15 @@
 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
 
 from airflow import settings
 from airflow._shared.timezones import timezone
 from airflow.configuration import conf
 from airflow.dag_processing.bundles.local import LocalDagBundle
+from airflow.dag_processing.importer_routing import get_claiming_coordinator
 from airflow.exceptions import (
     AirflowClusterPolicyError,
     AirflowClusterPolicySkipDag,
@@ -45,6 +46,7 @@
 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
@@ -58,6 +60,9 @@
     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):
     """
@@ -206,7 +211,7 @@ def __init__(
 
         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)
@@ -320,7 +325,7 @@ def _find_definitions(
 
     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
@@ -358,9 +363,11 @@ def _process_definition(
                 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:
@@ -431,36 +438,41 @@ def _get_relative_fileloc(self, filepath: str) -> str:
                 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()
+        if not isinstance(dag, SerializedLangSDKDAG):
+            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,
-                    )
-
-                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:
@@ -585,16 +597,30 @@ def sync_bag_to_db(
     version_data: dict[str, Any] | None = None,
     session: Session = NEW_SESSION,
 ) -> None:
-    """Save attributes about list of DAG to the DB."""
+    """
+    Save attributes about list of DAG to the DB.
+
+    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
 
-    import_errors = {(bundle_name, rel_path): error for rel_path, error in dagbag.import_errors.items()}
+    def is_parsed_by_runtime(rel_path: str) -> bool:
+        return get_claiming_coordinator(Path(dagbag.bundle_path or "", rel_path), bundle_name) is not None
+
+    import_errors = {
+        (bundle_name, rel_path): error
+        for rel_path, error in dagbag.import_errors.items()
+        if not is_parsed_by_runtime(rel_path)
+    }
 
     # Build the set of all files that were parsed and include files with import errors
     # in case they are not in parsed_definitions
     files_parsed = set(import_errors)
     if dagbag.bundle_path:
         for rel_path in dagbag.parsed_definitions:
+            if is_parsed_by_runtime(rel_path):
+                continue
             files_parsed.add((bundle_name, rel_path))
             # A definition nested in an archive also clears the archive's own discovery errors.
             if enclosing_file := find_enclosing_file(Path(dagbag.bundle_path, rel_path)):
@@ -603,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,

airflow-core/src/airflow/dag_processing/importer_routing.py

@@ -0,0 +1,59 @@
+#
+# 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.
+"""Route Dag files to the process that parses them, by the bundle's Dag importer registry."""
+
+from __future__ import annotations
+
+import logging
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter  # noqa: SDK001
+from airflow.sdk.importers import DagImporterRegistry, get_importer_registry  # noqa: SDK001
+
+if TYPE_CHECKING:
+    from airflow.sdk.coordinators._subprocess import SubprocessCoordinator  # noqa: SDK001
+
+log = logging.getLogger(__name__)
+
+
+def _get_registry(bundle_name: str | None) -> DagImporterRegistry | None:
+    try:
+        return get_importer_registry(bundle_name)
+    except Exception:
+        log.exception("Cannot build the Dag importer registry for bundle %s", bundle_name)
+        return None
+
+
+def get_claiming_coordinator(
+    path: str | os.PathLike[str], bundle_name: str | None
+) -> SubprocessCoordinator | None:
+    """
+    Return the coordinator whose runtime parses ``path``, or ``None`` when a Python child parses it.
+
+    A runtime parses the file when its importer is a coordinator's Dag importer.
+    """
+    if (registry := _get_registry(bundle_name)) is None:
+        return None
+    try:
+        importer = registry.get_importer(Path(path))
+    except Exception:
+        log.exception("Cannot load the Dag importer for %s", path)
+        return None
+    return importer.coordinator if isinstance(importer, CoordinatorDagImporter) else None