uranusjr commented on code in PR #74586: URL: https://github.com/apache/airflow/pull/74586#discussion_r4249125312
########## task-sdk/src/airflow/sdk/importers/yaml_importer/builder.py: ########## @@ -0,0 +1,254 @@ +# 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. + +""" +Build a validated :class:`~.models.DagDocument` into a :class:`~airflow.sdk.DAG`. + +The parser produces a validated document; this turns it into real DAG + operator +objects. Ready-made operators are instantiated from their ``uses`` import path; +tasks with ``run`` are built into stub operators and handled by AIP-108 task +handlers at execution time. + +Templating is explicit, realized importer-side (a la AIP-80). The format is +literal by default; only values the author marked render as templates. Value +grammar markers are only accepted on fields declared in ``template_fields``. + +Tasks are built in topological order so a task referenced by an upstream's XCom +already exists when the downstream operator is constructed. +""" + +from __future__ import annotations + +import keyword +import linecache +from typing import TYPE_CHECKING, Any, cast + +from airflow.sdk import DAG, XComArg +from airflow.sdk._shared.module_loading import import_string +from airflow.sdk.importers.yaml_importer.models import ( + ConstRef, + FileRef, + Literal, + TemplateRef, + TimetableSchedule, + XComRef, +) + +if TYPE_CHECKING: + from collections.abc import Callable + from pathlib import Path + + from airflow.sdk import BaseOperator + from airflow.sdk.importers.yaml_importer.models import DagDocument, Task, Value + + +class YamlDagBuildError(ValueError): + """A validated document could not be built into a DAG.""" + + +class DagBuilder: + """ + Build a :class:`DagDocument` into a :class:`DAG`. + + This holds the state a single build shares. The state belongs to one build + and is discarded with the instance. + """ + + def __init__(self, source_dir: Path | None) -> None: + self._source_dir = source_dir + self._operators: dict[str, BaseOperator] = {} + self._stub_callables: dict[tuple[str, ...], Callable[..., None]] = {} + + def build(self, doc: DagDocument) -> DAG: + tasks = {t.id_: t for t in doc.tasks} + with DAG(dag_id=doc.dag_id, schedule=_build_schedule(doc.schedule), **(doc.model_extra or {})) as dag: + # A task depended by another task must be built first. + for task_id in _topological_order(doc.tasks): + self._operators[task_id] = self._build_operator(tasks[task_id]) + + # Wire dependency edges explicitly. + for task in doc.tasks: + for upstream in sorted(task.get_dependency_ids()): + self._operators[upstream] >> self._operators[task.id_] + return dag + + def _build_stub_task(self, task: Task, args: dict[str, Value]) -> BaseOperator: + """ + Build a stub from a "run:" task. + + The task body is a registered SDK task handler resolved at run time; the + run mapping binds as that handler's arguments, where XCom and template + references resolve natively. The stub is fed a callable whose parameters + are the run argument names, so serialization can build and later + validate (AIP-108) the handler's argument spec. + """ + from airflow.sdk import task as task_decorators # Lazy to avoid provider registration. + + run_kwargs: dict[str, Any] = {} + for name, value in args.items(): + if not name.isidentifier() or keyword.iskeyword(name): + raise YamlDagBuildError(f"task {task.id_!r}: run argument {name!r} is not a valid identifier") + run_kwargs[name] = self._lower(value)[0] + decorator = task_decorators.stub(task_id=task.id_, **(task.model_extra or {})) + return decorator(self._stub_callable(tuple(run_kwargs)))(**run_kwargs).operator + + def _build_operator(self, task: Task) -> BaseOperator: + """Instantiate *task* as an operator.""" + if (args := task.run) is not None: + return self._build_stub_task(task, args) + + op_cls: Any + if task.uses is None: # uses: null (or omitted) becomes EmptyOperator. + from airflow.providers.standard.operators.empty import EmptyOperator as op_cls # noqa: SDK002 + else: + op_cls = import_string(task.uses) + + # Validate template fields. + template_fields = set(getattr(op_cls, "template_fields", ())) + kwargs: dict[str, Any] = {} + render_fields: list[str] = [] + off_template: list[str] = [] + for name, value in task.with_.items(): + kwargs[name], needs_render, has_marker = self._lower(value) + if has_marker and name not in template_fields: + off_template.append(name) + if needs_render: + render_fields.append(name) + if off_template: + raise YamlDagBuildError(f"task {task.id_!r}: value grammar only allowed on template_fields") + if extras := task.model_extra: # Merge in task-level arguments (raw, never templated). + kwargs.update(extras) + + return _apply_template_lowering(op_cls(task_id=task.id_, **kwargs), tuple(render_fields)) + + def _lower(self, value: Any) -> tuple[Any, bool, bool]: + """ + Lower a resolved value-grammar node. + + The return value is ``(operator argument, needs_render, has_marker)``. + + * ``needs_render`` is True if the node carries ``$t`` or ``$x``, and we + must validate it's in the operator's ``template_fields``. + reference anywhere, i.e. Review Comment: Fixed. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
