Skip to content

Commit 2b76bb3

Browse files
committed
TS SDK: cover native TypeScript Dags end to end
1 parent ae13e82 commit 2b76bb3

22 files changed

Lines changed: 1090 additions & 52 deletions

File tree

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -455,6 +455,12 @@ in the task spec.
455455
Templated arguments pass through untouched, so ``{{ ds }}`` in a ``conf`` value is rendered
456456
server-side where rendering already happens.
457457

458+
A worked example
459+
~~~~~~~~~~~~~~~~
460+
461+
``ts-sdk/example/src/native.ts`` puts the constructs above into one Dag, registered on the same
462+
bundle as the mixed-language handlers beside it, so a single artifact serves both authoring modes.
463+
458464
``new Dag`` and ``dag.task`` both take a trailing spec of Airflow options:
459465
``{ schedule: "@daily", tags: ["etl"] }`` for the Dag, ``{ retries: 2, retryDelay: 30 }`` for a task.
460466

‎airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import time
2020
from datetime import datetime, timezone
2121
from functools import cached_property
22+
from http import HTTPStatus
2223

2324
import boto3
2425
import requests
@@ -163,6 +164,21 @@ def get_variable(self, key: str):
163164
"""Get an Airflow Variable via API."""
164165
return self._make_request(method="GET", endpoint=f"variables/{key}")
165166

167+
def set_variable(self, key: str, value: str, description: str | None = None):
168+
"""Create or replace an Airflow Variable via API."""
169+
body = {"key": key, "value": value, "description": description}
170+
try:
171+
return self._make_request(method="POST", endpoint="variables", json=body)
172+
except requests.HTTPError as exc:
173+
# 409 == it already exists, from an earlier run of the same suite.
174+
if exc.response is None or exc.response.status_code != HTTPStatus.CONFLICT:
175+
raise
176+
return self._make_request(method="PATCH", endpoint=f"variables/{key}", json=body)
177+
178+
def get_tasks(self, dag_id: str):
179+
"""List a Dag's tasks, with the edges each one carries."""
180+
return self._make_request(method="GET", endpoint=f"dags/{dag_id}/tasks")
181+
166182
def trigger_dag_and_wait(self, dag_id: str, json=None):
167183
"""Trigger a DAG and wait for it to complete."""
168184
self.un_pause_dag(dag_id)
Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""
18+
End-to-end test of a Dag declared entirely in TypeScript.
19+
20+
Run with::
21+
22+
E2E_TEST_MODE=ts_sdk RUN_TS_SDK_NATIVE_DAG_TESTS=true \\
23+
uv run --project airflow-e2e-tests pytest \\
24+
tests/airflow_e2e_tests/ts_sdk_tests/test_ts_sdk_native_dag.py -xvs
25+
26+
Unlike ``test_ts_sdk_dag.py``, no Python file declares this Dag: the Dag processor asks the
27+
``airflow-ts-pack`` bundle to parse itself, and the bundle answers with the serialized Dag that
28+
``ts-sdk/example/src/native.ts`` built. The graph is a graph rather than a chain -- a task group, a
29+
named fan-in, order-only edges, a conditional, a multi-way branch, and a ``TriggerDagRunOperator`` a
30+
Python worker runs.
31+
32+
Gated behind ``RUN_TS_SDK_NATIVE_DAG_TESTS`` because it needs a Dag processor that can dispatch a
33+
parse request to a language coordinator. Until that lands, the bundle answers a request nothing
34+
sends, and the Dag never appears.
35+
"""
36+
37+
from __future__ import annotations
38+
39+
import os
40+
from dataclasses import dataclass
41+
from datetime import datetime, timezone
42+
43+
import pytest
44+
45+
from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
46+
47+
_RUN_NATIVE = os.environ.get("RUN_TS_SDK_NATIVE_DAG_TESTS", "").lower() in ("true", "1")
48+
49+
# Parsing the bundle launches node before the first task is even scheduled, so
50+
# allow the same headroom the mixed-language suite does.
51+
_TS_TASK_TIMEOUT = 600
52+
53+
_DAG_ID = "typescript_native_example"
54+
55+
# Read by the handlers; see ts-sdk/example/src/native.ts.
56+
_NORTH_ROWS_VARIABLE = "typescript_native_north_rows"
57+
_SOUTH_ROWS_VARIABLE = "typescript_native_south_rows"
58+
_CADENCE_VARIABLE = "typescript_native_cadence"
59+
60+
pytestmark = pytest.mark.skipif(
61+
not _RUN_NATIVE,
62+
reason="Needs a Dag processor that dispatches parse requests to a language coordinator "
63+
"(RUN_TS_SDK_NATIVE_DAG_TESTS)",
64+
)
65+
66+
67+
@dataclass
68+
class _CompletedRun:
69+
client: AirflowClient
70+
run_id: str
71+
state: str
72+
ti_states: dict[str, str]
73+
74+
def xcom(self, task_id: str, key: str = "return_value"):
75+
return self.client.get_xcom_value(dag_id=_DAG_ID, task_id=task_id, run_id=self.run_id, key=key).get(
76+
"value"
77+
)
78+
79+
80+
@pytest.fixture(scope="module")
81+
def parsed_dag() -> AirflowClient:
82+
"""A client that has waited for the bundle's Dag to be parsed and registered."""
83+
client = AirflowClient()
84+
# The Dag processor spawns node to parse the bundle, so the Dag appears some
85+
# time after the deployment is up; every read below would 404 until then.
86+
client.wait_for_dag(_DAG_ID, timeout=_TS_TASK_TIMEOUT)
87+
return client
88+
89+
90+
@pytest.fixture(scope="module")
91+
def completed_run(parsed_dag: AirflowClient) -> _CompletedRun:
92+
"""Trigger the native Dag once, with the inputs every test below reads back."""
93+
client = parsed_dag
94+
# Both regions non-empty, so the conditional takes its `then` branch, and a
95+
# weekly cadence so the branch's choice is known rather than guessed.
96+
for key, value in (
97+
(_NORTH_ROWS_VARIABLE, "3"),
98+
(_SOUTH_ROWS_VARIABLE, "2"),
99+
(_CADENCE_VARIABLE, "weekly"),
100+
):
101+
client.set_variable(key, value)
102+
103+
client.un_pause_dag(_DAG_ID)
104+
resp = client.trigger_dag(_DAG_ID, json={"logical_date": datetime.now(timezone.utc).isoformat()})
105+
run_id = resp["dag_run_id"]
106+
state = client.wait_for_dag_run(dag_id=_DAG_ID, run_id=run_id, timeout=_TS_TASK_TIMEOUT)
107+
ti_resp = client.get_task_instances(dag_id=_DAG_ID, run_id=run_id)
108+
return _CompletedRun(
109+
client=client,
110+
run_id=run_id,
111+
state=state,
112+
ti_states={ti["task_id"]: ti.get("state") for ti in ti_resp.get("task_instances", [])},
113+
)
114+
115+
116+
def test_the_dag_the_bundle_parsed_is_registered(parsed_dag: AirflowClient):
117+
"""The Dag exists without any Python file declaring it."""
118+
dag = parsed_dag.get_dag(_DAG_ID)
119+
120+
assert dag["dag_id"] == _DAG_ID
121+
# The serializer expands a cron preset, so "@daily" is recorded as its expression.
122+
assert dag.get("timetable_summary") == "0 0 * * *"
123+
# `tags` is a list of objects, each naming one tag.
124+
assert {tag["name"] for tag in dag.get("tags") or []} >= {"typescript", "native"}
125+
126+
127+
def test_the_graph_carries_every_construct(parsed_dag: AirflowClient):
128+
"""Group prefixes, the fan-in, the branches and the trigger all survive parsing."""
129+
tasks = parsed_dag.get_tasks(_DAG_ID).get("tasks", [])
130+
downstream = {task["task_id"]: set(task["downstream_task_ids"]) for task in tasks}
131+
132+
# The task group prefixed its members.
133+
assert {"extract.north", "extract.south"} <= set(downstream)
134+
# The named fan-in put both extract tasks upstream of summarize. The API
135+
# reports downstream edges only, so upstream is read by inverting them.
136+
assert {"extract.north", "extract.south"} <= {
137+
task_id for task_id, down in downstream.items() if "summarize" in down
138+
}
139+
# The conditional and the branch reach their candidates.
140+
assert downstream["has_rows"] >= {"load_rows", "report_empty"}
141+
assert downstream["pick_cadence"] >= {"publish_daily", "publish_weekly"}
142+
# The group edge, expanded onto the tasks the group leaves from.
143+
assert {"extract.north", "extract.south"} <= {
144+
task_id for task_id, down in downstream.items() if "pick_cadence" in down
145+
}
146+
# Cleanup sits behind every branch outcome, which is what its trigger rule is for.
147+
assert {"load_rows", "report_empty", "publish_daily", "publish_weekly"} <= {
148+
task_id for task_id, down in downstream.items() if "cleanup" in down
149+
}
150+
# And the one task Python runs rather than the Node coordinator.
151+
by_id = {task["task_id"]: task for task in tasks}
152+
assert by_id["trigger_downstream"]["operator_name"] == "TriggerDagRunOperator"
153+
154+
155+
def test_dag_run_succeeded(completed_run: _CompletedRun):
156+
assert completed_run.state == "success", (
157+
f"expected the run to succeed; got {completed_run.state!r}. task states: {completed_run.ti_states}"
158+
)
159+
160+
161+
def test_the_taken_branches_ran_and_the_others_skipped(completed_run: _CompletedRun):
162+
"""A branch is a run-time skip, so the states are what prove it worked."""
163+
states = completed_run.ti_states
164+
165+
# Both regions had rows, so the conditional followed `then`.
166+
assert states.get("load_rows") == "success"
167+
assert states.get("report_empty") == "skipped"
168+
169+
# The cadence Variable said weekly, so that is the case the branch chose.
170+
assert states.get("publish_weekly") == "success"
171+
assert states.get("publish_daily") == "skipped"
172+
173+
174+
def test_every_other_task_succeeded(completed_run: _CompletedRun):
175+
always_run = [
176+
"extract.north",
177+
"extract.south",
178+
"summarize",
179+
"has_rows",
180+
"pick_cadence",
181+
"cleanup",
182+
"trigger_downstream",
183+
]
184+
for task_id in always_run:
185+
assert completed_run.ti_states.get(task_id) == "success", (
186+
f"{task_id!r} did not succeed. all task states: {completed_run.ti_states}"
187+
)
188+
189+
190+
def test_xcoms_flow_between_typescript_tasks(completed_run: _CompletedRun):
191+
"""The fan-in's arguments reached the handler, and its own push came back."""
192+
assert completed_run.xcom("extract.north") == {"region": "north", "rows": 3}
193+
assert completed_run.xcom("extract.south") == {"region": "south", "rows": 2}
194+
assert completed_run.xcom("summarize") == {"total": 5, "regions": 2}
195+
# Pushed under its own key by the summarize handler, and read back by load_rows.
196+
assert completed_run.xcom("summarize", key="region_total") == 5
197+
assert completed_run.xcom("load_rows") == {"loaded": 5}
198+
199+
200+
def test_the_decisions_are_recorded(completed_run: _CompletedRun):
201+
"""A condition returns its boolean, and a branch the task id it chose."""
202+
assert completed_run.xcom("has_rows") is True
203+
assert completed_run.xcom("pick_cadence") == "publish_weekly"
204+
205+
206+
def test_a_cleared_skipped_branch_stays_skipped(completed_run: _CompletedRun):
207+
"""What `_can_skip_downstream` plus the skipmixin XCom are for."""
208+
assert completed_run.xcom("has_rows", key="skipmixin_key") == {"skipped": ["report_empty"]}
209+
assert completed_run.xcom("pick_cadence", key="skipmixin_key") == {"skipped": ["publish_daily"]}

0 commit comments

Comments
 (0)