FrankYang0529 commented on code in PR #75190:
URL: https://github.com/apache/airflow/pull/75190#discussion_r4344185315
##########
ts-sdk/src/sdk/dag.ts:
##########
@@ -327,48 +327,105 @@ export function isTaskRef(value: unknown): value is
TaskRef {
const conditionTasks = new WeakMap<object, TaskRef>();
-const labelledNodes = new WeakMap<object, { readonly node: Node; readonly
text: string }>();
+const labelledNodes = new WeakMap<
+ object,
+ { readonly nodes: readonly Node[]; readonly text: string }
+>();
function resolveNode(node: Node): Node {
if (typeof node !== "object" || node === null) return node;
- return labelledNodes.get(node)?.node ?? conditionTasks.get(node) ?? node;
+ return conditionTasks.get(node) ?? node;
}
/**
- * Label the edge that `before` or `after` draws to `node`, as Python's
`Label` does.
- * The result only goes into `before` or `after`; calling either on it throws.
- * `text` cannot be empty.
+ * A text label for an edge. The Airflow UI shows the label on the edge in the
graph, as it does
+ * for Python's `Label`. The text comes before the task that the label is put
on. That way the
+ * text sits between the two tasks of the edge in the code, just as the label
sits between them in
+ * the graph:
*
* ```ts
- * checked.before(label(processed, "rows found"), label(emptyNotice, "no
rows"));
+ * // load >> Label("when empty") >> report_empty
+ * loaded.before(Label.of("when empty").on(reportedEmpty));
+ * // load >> Label("always") >> cleanup
+ * cleaned.after(Label.of("always").on(loaded));
* ```
*/
-export function label(node: Node, text: string): Node {
- if (typeof text !== "string") {
- throw new Error(`label() takes a string to name the edge with, not
${describeValue(text)}`);
+export class Label {
+ readonly #text: string;
+
+ // The constructor is private only in TypeScript. Plain JavaScript can still
call
+ // `new Label("")`, so the text check lives here and not in `of`. A check in
`of` would not run
+ // for that call.
+ private constructor(text: string) {
+ if (typeof text !== "string") {
+ throw new Error(
+ `Label.of() takes a string to name the edge with, not
${describeValue(text)}`,
+ );
+ }
+ if (text.trim() === "") {
+ throw new Error(
+ "Label.of() takes the text to put on the edge, not " +
+ (text === "" ? "an empty string" : "a blank string"),
+ );
+ }
+ this.#text = text;
}
- if (text === "") {
- throw new Error("label() takes the text to put on the edge, not an empty
string");
+
+ /**
+ * Make a label with the given `text`. Use `on` to put the label on edges.
`text` must not be
+ * empty or only whitespace.
+ */
+ static of(text: string): Label {
+ return new Label(text);
}
- const target = resolveNode(node);
- const id = nodeId(target);
- if (id === undefined) {
- throw new Error(`label() takes a task or task group reference, not
${describeValue(node)}`);
+
+ /**
+ * Put the label on each edge that a `before` or `after` call draws between
its receiver and
+ * `node` or a node in `more`.
+ *
+ * So each edge of a fan-out can have its own label, and one label can go on
several edges:
+ *
+ * ```ts
+ * checked.before(Label.of("rows found").on(processed), Label.of("no
rows").on(reportedEmpty));
+ * // load >> Label("forks") >> [report_empty, notify]
+ * loaded.before(Label.of("forks").on(reportedEmpty, notified));
+ * ```
+ *
+ * Pass the value that `on` returns to `before` or `after`. Calling `before`
or `after` on that
+ * value throws an error.
+ */
+ on(node: Node, ...more: readonly Node[]): Node {
+ const text = this.#text;
+ const nodes = [node, ...more].flatMap((each) => {
+ // Labelling a labelled node again replaces its label.
+ const labelled = labelledNodes.get(each);
+ if (labelled !== undefined) return labelled.nodes;
+ const target = resolveNode(each);
+ if (nodeId(target) === undefined) {
+ throw new Error(
+ `Label.of(${JSON.stringify(text)}).on() takes a task or task group
reference, ` +
+ `not ${describeValue(each)}`,
+ );
+ }
+ return [target];
+ });
+ const { dagId } = nodes[0]!;
Review Comment:
`on()` now throws when its nodes belong to different Dags, before any edge
is drawn.
--
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]
