Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -365,6 +365,21 @@ repos:
^scripts/ci/prek/check_go_sdk_spec_drift\.py$
pass_filenames: false
require_serial: true
# 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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,29 @@ Pass ``{ prefixGroupId: false }`` to keep the ids declared in a group as written
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" });
Comment thread
pierrejeambrun marked this conversation as resolved.

// ...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.

``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.

Expand Down
16 changes: 16 additions & 0 deletions scripts/ci/lang_sdk_serialization/__init__.py
Original file line number Diff line number Diff line change
@@ -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.
233 changes: 233 additions & 0 deletions scripts/ci/lang_sdk_serialization/compare.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,233 @@
# 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. The Python side also takes the SDK's output as Airflow
receives it, filling in the fields the SDK leaves to Airflow's config, and loads it through Airflow's
deserializer. The two are then compared field by field, so this fails both when the 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 the 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",
Comment thread
guan404ming marked this conversation as resolved.
}
)


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"
received_output = directory / f"received_{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),
"--receive",
str(sdk_output),
str(received_output),
]
)

python = json.loads(python_output.read_text())
problems = compare(python, json.loads(received_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"The 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())
Loading