TS SDK: add conditional branching with dag.if

.pre-commit-config.yaml

@@ -337,6 +337,21 @@ repos:
           (?x)
           ^airflow-core/src/airflow/serialization/schema\.json$|
           ^ts-sdk/schema/dag-schema\.json$
+      # check-ts-sdk-serialization-conformance in ts-sdk/ covers the SDK's own files; this runs the same
+      # check when the shared harness or Airflow's serializer changes, which that project cannot see.
+      - id: check-ts-sdk-serialization-conformance-shared
+        name: Check the TS SDK serializes Dags the way Airflow does, after a shared change
+        description: "Serialize the shared test Dags with the TS SDK and with Airflow, and compare the two"
+        entry: ./ts-sdk/scripts/ci/prek/check_serialization_conformance.py
+        language: node
+        additional_dependencies: ['pnpm@10.28.1']
+        pass_filenames: false
+        require_serial: true
+        files: >
+          (?x)
+          ^airflow-core/src/airflow/serialization/serialized_objects\.py$|
+          ^airflow-core/src/airflow/serialization/schema\.json$|
+          ^scripts/ci/lang_sdk_serialization/.*$
       - id: check-go-version-in-sync
         name: Check Go toolchain version is consistent across build files
         entry: ./scripts/ci/prek/check_go_version_in_sync.py

airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst

@@ -311,6 +311,78 @@ already exists changes nothing. Each returns the reference it was called on, so
 Pass a value as an argument when the downstream task needs it, and use ``before`` or ``after`` when
 it only needs to run in order.
 
+Task groups
+~~~~~~~~~~~
+
+``dag.taskGroup(groupId)`` opens a scope with the same ``task`` and ``taskGroup`` methods as the Dag,
+prefixing the id of everything declared in it, as Python's ``prefix_group_id`` does:
+
+.. code-block:: typescript
+
+    const staging = dag.taskGroup("staging");
+    staging.task("stage_rows", stageRows)();        // task id "staging.stage_rows"
+    staging.taskGroup("checks").task("nulls", checkNulls)();  // "staging.checks.nulls"
+
+    staging.before(loaded);                          // staging >> loaded
+
+A group is an edge endpoint in its own right, so ``before`` and ``after`` order a whole group against
+a task or against another group.
+
+Tasks and groups share one id namespace, as they do in Python, so a Dag cannot hold both a task and a
+group called ``staging``. A ``.`` is what separates a group from what it holds, so it cannot appear in
+an id of either.
+
+Pass ``{ prefixGroupId: false }`` to keep the ids declared in a group as written, as ``prefix_group_id=False``
+does in Python; they then have to be unique across the Dag. A group id is made of letters, digits, dashes and
+underscores, and is at most 200 characters.
+
+Serialization
+~~~~~~~~~~~~~
+
+A native Dag serializes into the same Dag JSON a Python Dag produces, so the scheduler reads it
+without knowing which language declared it.
+
+``schedule`` accepts what maps to a stock timetable: unset, ``@once``, ``@continuous``, or a cron
+expression. A cron preset such as ``@daily`` is recorded as the expression it stands for. Anything
+else names a Python object a TypeScript bundle cannot point at, and is rejected.
+
+Every task of a native Dag runs on the Node coordinator, so it needs the queue the deployment routes
+there. Set it once on the Dag and each task inherits it:
+
+.. code-block:: typescript
+
+    const dag = new Dag("ts_etl", { schedule: "@daily", queue: "typescript" });
+
+    // ...and one task that needs its own.
+    dag.task("heavy", heavyHandler, { queue: "typescript_large" })();
+
+``queue`` on a task wins over the Dag's. See :ref:`typescript-sdk/coordinator-config` for the
+``queue_to_coordinator`` entry that sends that queue to the coordinator.
+
+Conditional branching
+~~~~~~~~~~~~~~~~~~~~~
+
+``dag.if`` takes a task whose handler returns a boolean, and names the task each outcome runs:
+
+.. code-block:: typescript
+
+    const condition = dag.task("has_rows", async ({ rows }: { rows: number }) => rows > 0);
+    const gated = condition({ rows: extracted });
+
+    dag.if(gated).then(loaded).else(reportedEmpty);
+
+The condition is an ordinary task, so it is declared, typed and wired like any other, and the
+compiler checks that its handler really returns a boolean. ``else`` is optional: a one-sided
+condition skips its own branch when the condition fails and follows nothing.
+
+A guarded task takes no argument for the control edge, because a condition's boolean decides whether
+the task runs rather than what it runs on. Read a value from the condition with
+``getClient().getXCom``.
+
+The side not taken is skipped when the run reaches it, and stays skipped if you clear it later. Only
+the branches named here are skipped, so a task that several branches converge on still runs — unlike
+Python's ``@task.branch``, which skips every immediate downstream it did not follow.
+
 ``new Dag`` and ``dag.task`` both take a trailing spec of Airflow options:
 ``{ schedule: "@daily", tags: ["etl"] }`` for the Dag, ``{ retries: 2, retryDelay: 30 }`` for a task.
 

scripts/ci/lang_sdk_serialization/__init__.py

@@ -0,0 +1,16 @@
+# 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.

scripts/ci/lang_sdk_serialization/compare.py

@@ -0,0 +1,230 @@
+# 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.
+r"""
+Check that a language SDK serializes Dags the way Airflow does.
+
+Runs the SDK's serializer and serialize_python.py over test_dags.yaml, each writing its serialization
+to a JSON file in a temporary directory, and compares the two field by field. The Python side also
+loads the SDK's output through Airflow's deserializer, so this fails both when the two serializers
+drift apart and when Airflow cannot read what the SDK writes.
+
+The SDK's serializer is the command after ``--``. It runs from the repository root, with the paths of
+test_dags.yaml and of the JSON file to write appended, and writes each Dag as
+``DagSerialization.to_dict`` would, keyed by Dag id. The Python side needs ``uv``::
+
+    python3 scripts/ci/lang_sdk_serialization/compare.py --sdk typescript -- \
+        pnpm --dir ts-sdk exec tsx tests/conformance/serialize_typescript.ts
+
+On a failure both files are kept, and their directory is printed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any
+
+HERE = Path(__file__).resolve().parent
+REPO_ROOT = HERE.parents[2]
+TEST_DAGS = HERE / "test_dags.yaml"
+SCHEMA = REPO_ROOT / "airflow-core" / "src" / "airflow" / "serialization" / "schema.json"
+
+# Name the file a Dag was declared in: Python's own Dag file, or the SDK's bundle.
+DAG_KEYS_NOT_COMPARED = frozenset({"fileloc", "relative_fileloc", "_processor_dags_folder"})
+
+TASK_KEYS_NOT_COMPARED = frozenset(
+    {
+        # Operator identity: Python names the Python class that ran, a language SDK a fixed pair.
+        # Neither is ever imported on the Airflow side.
+        "task_type",
+        "_task_module",
+        # The SDK's marker for a task in its language, which Python has no equivalent for.
+        "language",
+        # Python bookkeeping for mapped tasks and retry policies, neither of which a language SDK's
+        # Dag can declare.
+        "_needs_expansion",
+        "has_retry_policy",
+        # The stub contract every language SDK task carries, which the Python operator built here
+        # has no counterpart for. Loading the SDK's output through Airflow checks them.
+        "is_stub",
+        "_arg_bindings",
+    }
+)
+
+
+def run(command: list[str]) -> None:
+    if subprocess.run(command, cwd=REPO_ROOT, check=False).returncode:
+        raise SystemExit(f"`{' '.join(command)}` failed")
+
+
+def get_task_defaults() -> dict[str, Any]:
+    """Map each task field the Dag schema gives a default to that default."""
+    fields = json.loads(SCHEMA.read_text())["definitions"]["operator"]["properties"]
+    return {key: field["default"] for key, field in fields.items() if field.get("default") is not None}
+
+
+def normalize_as_javascript(value: Any) -> Any:
+    """Read a JSON value as JavaScript does: one number type, and a bool that is not a number."""
+    if isinstance(value, bool):
+        return ("bool", value)
+    if isinstance(value, (int, float)):
+        return float(value)
+    if isinstance(value, list):
+        return [normalize_as_javascript(item) for item in value]
+    if isinstance(value, dict):
+        return {key: normalize_as_javascript(item) for key, item in value.items()}
+    return value
+
+
+def is_same_json(python: Any, sdk: Any) -> bool:
+    return normalize_as_javascript(python) == normalize_as_javascript(sdk)
+
+
+def find_differences(path: str, python: Any, sdk: Any) -> list[str]:
+    """List where two JSON values differ, down to the innermost key or index."""
+    if isinstance(python, dict) and isinstance(sdk, dict):
+        problems = []
+        for key in sorted(python.keys() | sdk.keys()):
+            if key not in sdk:
+                problems.append(f"{path}.{key} is missing, Python writes {python[key]!r}")
+            elif key not in python:
+                problems.append(f"{path}.{key} is {sdk[key]!r}, which Python does not write")
+            else:
+                problems.extend(find_differences(f"{path}.{key}", python[key], sdk[key]))
+        return problems
+    if isinstance(python, list) and isinstance(sdk, list) and len(python) == len(sdk):
+        return [
+            problem
+            for index, (python_item, sdk_item) in enumerate(zip(python, sdk))
+            for problem in find_differences(f"{path}[{index}]", python_item, sdk_item)
+        ]
+    if not is_same_json(python, sdk):
+        return [f"{path} is {sdk!r}, Python writes {python!r}"]
+    return []
+
+
+def compare_fields(
+    python: dict[str, Any], sdk: dict[str, Any], not_compared: frozenset[str], defaults: dict[str, Any]
+) -> list[str]:
+    """
+    Compare two serialized objects key by key.
+
+    A key the SDK leaves out is fine when Python wrote its schema default. Python keeps such a value
+    when its ``client_defaults`` table disagrees with the schema, a table a language SDK is never sent,
+    and Airflow reads a missing field as its default anyway.
+    """
+    problems = []
+    for key in sorted((python.keys() | sdk.keys()) - not_compared):
+        if key in python and key in sdk:
+            problems.extend(find_differences(key, python[key], sdk[key]))
+        elif key in sdk:
+            problems.append(f"{key} is {sdk[key]!r}, which Python does not write")
+        elif key not in defaults or not is_same_json(python[key], defaults[key]):
+            problems.append(f"{key} is missing, Python writes {python[key]!r}")
+    return problems
+
+
+def compare_dag(python: dict[str, Any], sdk: dict[str, Any], defaults: dict[str, Any]) -> list[str]:
+    problems = compare_fields(python, sdk, DAG_KEYS_NOT_COMPARED | {"tasks"}, {})
+    python_ids = [task["__var"]["task_id"] for task in python["tasks"]]
+    sdk_ids = [task["__var"]["task_id"] for task in sdk["tasks"]]
+    if sdk_ids != python_ids:
+        return [*problems, f"tasks are {sdk_ids}, Python writes {python_ids}"]
+    for python_task, sdk_task in zip(python["tasks"], sdk["tasks"]):
+        task_id = python_task["__var"]["task_id"]
+        if sdk_task["__type"] != python_task["__type"]:
+            problems.append(f"task {task_id} is a {sdk_task['__type']!r}, not a {python_task['__type']!r}")
+        problems.extend(
+            f"task {task_id}: {problem}"
+            for problem in compare_fields(
+                python_task["__var"], sdk_task["__var"], TASK_KEYS_NOT_COMPARED, defaults
+            )
+        )
+    return problems
+
+
+def compare(python: dict[str, Any], sdk: dict[str, Any], task_defaults: dict[str, Any]) -> list[str]:
+    """List every way the SDK's serialization differs from Python's."""
+    if sdk.keys() != python.keys():
+        return [f"the Dags are {sorted(sdk)}, Python writes {sorted(python)}"]
+    problems = []
+    for dag_id, python_dag in python.items():
+        sdk_dag = sdk[dag_id]
+        if sdk_dag["__version"] != python_dag["__version"]:
+            problems.append(
+                f"{dag_id}: __version is {sdk_dag['__version']}, Python writes {python_dag['__version']}"
+            )
+        problems.extend(
+            f"{dag_id}: {problem}"
+            for problem in compare_dag(python_dag["dag"], sdk_dag["dag"], task_defaults)
+        )
+    return problems
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--sdk", required=True, help="the SDK's name, as in serialized_<sdk>.json")
+    parser.add_argument("command", nargs="+", help="the SDK's serializer, after --")
+    args = parser.parse_args(argv)
+
+    directory = Path(tempfile.mkdtemp(prefix=f"{args.sdk}-serialization-"))
+    python_output = directory / "serialized_python.json"
+    sdk_output = directory / f"serialized_{args.sdk}.json"
+    run([*args.command, str(TEST_DAGS), str(sdk_output)])
+    run(
+        [
+            "uv",
+            "run",
+            "--project",
+            "airflow-core",
+            # airflow-core's dev group pulls in providers and extras that build native code; the
+            # serializer only needs airflow-core itself.
+            "--no-dev",
+            "python",
+            str(HERE / "serialize_python.py"),
+            str(TEST_DAGS),
+            str(python_output),
+            "--deserialize",
+            str(sdk_output),
+        ]
+    )
+
+    python = json.loads(python_output.read_text())
+    problems = compare(python, json.loads(sdk_output.read_text()), get_task_defaults())
+    if problems:
+        print(
+            f"The {args.sdk} serialization differs from Python's in {len(problems)} place(s):",
+            file=sys.stderr,
+        )
+        for problem in problems:
+            print(f"  {problem}", file=sys.stderr)
+        print(f"Both serializations are kept in {directory}", file=sys.stderr)
+        return 1
+    shutil.rmtree(directory)
+    print(f"The {args.sdk} SDK serializes all {len(python)} Dags of {TEST_DAGS.name} as Python does")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())