kaxil commented on code in PR #74439: URL: https://github.com/apache/airflow/pull/74439#discussion_r4216808843
########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt: ########## @@ -0,0 +1,291 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk + +import kotlin.Throws + +/** + * The Airflow task ID of one task of a Dag. + * + * A `@Builder.Branch` method names the task it chose with one of these. The + * generated `<Dag>Builder.TaskIds` holds a constant per task of the Dag, so + * the choice is checked where it is written rather than when the task runs. + * + * @property value The task ID. + */ +class TaskId private constructor( + val value: String, +) { + companion object { + /** + * @suppress + * + * The ID [value] names. Public so a generated `TaskIds` holder can declare + * its constants; user code names a task through those constants. + */ + @JvmStatic + fun of(value: String): TaskId = TaskId(value) + } + + override fun equals(other: Any?): Boolean = other is TaskId && other.value == value + + override fun hashCode(): Int = value.hashCode() + + override fun toString(): String = value +} + +/** + * A task that chooses one of several tasks to run; every other one is skipped. + * + * Register one with [DagDef.Branch], then list what it can choose with + * [BranchRef.option]. The choice is the task's own class, so javac checks it: + * + * ```java + * public class PickPath implements BranchTask { + * @Override + * public Class<? extends Task> choose(Context context, Client client) { + * return ((Number) client.getXCom("extract")).longValue() > 1000 + * ? HandleLong.class + * : HandleShort.class; + * } + * } + * ``` + * + * The SDK runs [choose] and pushes the chosen task's ID as this task's return + * value, so [execute] is never called. + * + * Because a branch names a case by its class, no two of its cases may be + * registered from the same class; registering the Dag reports that. + * + * @see DagDef.Branch + */ +interface BranchTask : Task { + /** + * Chooses the one case that runs. + * + * Any exception thrown marks the task instance as failed, and nothing is + * then pushed or skipped. + * + * @param context Runtime context for the current execution workload. + * @param client Client for Airflow API calls scoped to this execution. + * @return The class of one of this branch's cases. + * @throws Exception on failure; the task instance is marked failed. + */ + @Throws(Exception::class) + fun choose( + context: Context, + client: Client, + ): Class<out Task> + + /** Never called: the SDK runs a branch through [choose]. */ + override fun execute( + context: Context, + client: Client, + ): Unit = + throw IllegalStateException( + "Branch '${javaClass.name}' runs through choose(), so execute() is never called", + ) +} + +/** + * @suppress + * + * What a generated `@Builder.Branch` task implements: a branch whose method + * named the case with a [TaskId]. Public so generated code can implement it; + * a Dag written against [DagDef.Branch] implements [BranchTask] instead. + */ +interface TaskIdBranchTask : Task { + @Throws(Exception::class) + fun choose( + context: Context, + client: Client, + ): TaskId + + override fun execute( + context: Context, + client: Client, + ): Unit = + throw IllegalStateException( + "Branch '${javaClass.name}' runs through choose(), so execute() is never called", + ) +} + +/** + * A branch registered with a Dag: list the tasks it can choose between. + * + * ```java + * dag.Branch(PickPath.class).option(handleLong).option(handleShort); + * ``` + * + * A case runs after the branch, so naming it records that edge, as + * [Deps.Flow.before] would. The branch chooses exactly one case and skips + * every other, so a task that runs after several cases needs a trigger rule + * that tolerates a skipped upstream, such as `none_failed_min_one_success`. + * + * @see DagDef.Branch + */ +class BranchRef private constructor( + private val ref: TaskRef<*>, + private val cases: BranchDef, +) : Deps.Flow { + companion object { + /** + * @suppress + * + * Marks a registered task as a branch. Public so a generated wiring view + * can call it; user code reaches a branch through [DagDef.Branch]. + */ + @JvmStatic + fun of(ref: TaskRef<*>): BranchRef { + require(ref.def.decider == null) { + "Task '${ref.def.id}' already decides what to skip; declare it once" + } + val definition = ref.def.definition + require( + BranchTask::class.java.isAssignableFrom(definition) || + TaskIdBranchTask::class.java.isAssignableFrom(definition), + ) { + "Task '${ref.def.id}' runs '${definition.name}', which chooses nothing; a branch runs a " + + "BranchTask" + } + return BranchRef(ref, BranchDef().also { ref.def.decider = it }) + } + } + + /** Airflow task ID of the deciding task. */ + val id: String get() = ref.def.id + + /** + * Adds a task this branch can choose. Every other case is skipped when the + * branch chooses this one. + * + * @param task Task of the same Dag. + * @return This branch, for chaining. + * @throws IllegalArgumentException if the task belongs to another Dag or is + * already a case of this branch. + */ + fun option(task: TaskRef<*>): BranchRef { + cases.add(ref.def, task.def) + // The case runs after the branch, which is what puts the branch in the + // serialized Dag as its upstream and lets registration see a cycle that + // runs through a branch. + ref.before(task) + return this + } + + /** + * Sets one task-level configuration value on the deciding task. + * + * @param key Airflow task setting name. + * @param value Value matching the key's schema type. + * @return This branch, for chaining. + * @throws IllegalArgumentException if the key is unknown or the value type + * does not match. + */ + fun config( + key: String, + value: Any?, + ): BranchRef { + ref.config(key, value) + return this + } + + override fun nodes(): List<TaskDef> = ref.nodes() + + override fun before(vararg next: Deps.Flow): BranchRef { + ref.before(*next) + return this + } + + override fun after(vararg previous: Deps.Flow): BranchRef { + ref.after(*previous) + return this + } +} + +/** The cases of one branch, as [BranchRef.option] listed them. */ +internal class BranchDef : DeciderDef { + private val options = mutableListOf<TaskDef>() + + override val cases: List<TaskDef> get() = options + + fun add( + decider: TaskDef, + case: TaskDef, + ) { + requireUnregistered(decider, "Branch") + requireSameDag(decider, case, "Branch") + require(options.none { it === case }) { + "Branch '${decider.id}' already chooses between '${case.id}' and others; name each case once" + } + // A BranchTask names its case by class, so two cases sharing one would be + // indistinguishable. Checked here rather than at registration so the error + // points at the second option(...) call. + if (BranchTask::class.java.isAssignableFrom(decider.definition)) { Review Comment: The conformance harness already hits this limit (`ConformanceBranch` falls back to the hidden `TaskIdBranchTask` because every task there shares one class), and a real Dag that registers one generic class under several ids, say `load_eu` and `load_us` both on `Load.class`, can't branch between them with the public API at all. Going the other way, the annotation style picks by id because there is no class to name, but the processor already generates one uniquely named class per task method (`AnnotationExampleBuilder.ReportLong`), so I think `byClass` would match that just as well. Would one mechanism for both styles be simpler? Either an id-based choice the interface style can use too, or class-based everywhere and drop `TaskId` and the generated `TaskIds`. As it stands, `TaskId.of(String)` is public too, which reopens the string escape hatch that ADR 0008 decision 2 closes. ########## java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt: ########## @@ -567,6 +599,29 @@ class BuilderProcessor : AbstractProcessor() { "Dag has both a task and a task group with ID '${group.fullId}'; rename one" } } + if (declarations.any { it.kind == TaskKind.BRANCH }) { + val byConstant = mutableMapOf<String, TaskDeclaration>() + declarations.forEach { decl -> + val constant = constantName(decl.id) + require(SourceVersion.isName(constant)) { Review Comment: This runs over every task in the Dag, not only the cases. So a Dag with a `2nd_pass` task (or both `load-data` and `load_data`) compiles today and stops compiling once someone adds a `@Builder.Branch` anywhere in it, and the fix the message suggests is renaming a task id, which leaves that task's run history behind under the old id. Could a constant that starts with a digit get a `_` prefix instead? For a task inside a group whose id starts with a digit, the suggested `id = "..."` can't help either, since the digit comes from the group. ########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt: ########## @@ -201,6 +201,47 @@ class DagDef( definition: Class<out ConditionTask>, ): ConditionRef = ConditionRef.of(task(id, definition)) + /** + * Declares a task that chooses one of several tasks to run; every other one + * is skipped. + * + * The task's ID is the class's simple name with its first character + * lowercased, so `PickPath.class` becomes `pickPath`. Use + * [Branch(id, definition)][Branch] to set it, and [BranchRef.config] for the + * task's other settings: + * + * ```java + * dag.Branch(PickPath.class).option(handleLong).option(handleShort); + * ``` + * + * @param definition Class that implements [BranchTask]. Must have a public + * no-arg constructor. + * @return The branch, to list its cases on. + * @throws IllegalArgumentException if a task with the same ID is already + * registered. + */ + @Suppress("ktlint:standard:function-naming") + fun Branch(definition: Class<out BranchTask>): BranchRef = Branch(deriveTaskId(definition), definition) Review Comment: ADR 0008 decision 1 names this construct `Switch`, and Go (`dag.Switch(...).Case(...)`) and TS (`dag.switch(...).case(...)`) both follow it. The conformance Dag this PR adds to the check is even called `switch`. `Switch` is as legal a Java method name as `If`. Is there a reason Java should be `Branch`/`option`? If so, it's probably worth recording in the ADR, since someone moving between SDKs will look for `Switch`. ########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt: ########## @@ -0,0 +1,291 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk + +import kotlin.Throws + +/** + * The Airflow task ID of one task of a Dag. + * + * A `@Builder.Branch` method names the task it chose with one of these. The + * generated `<Dag>Builder.TaskIds` holds a constant per task of the Dag, so + * the choice is checked where it is written rather than when the task runs. Review Comment: The generated `TaskIds` Javadoc says the opposite: every task of the Dag has a constant, so naming one that isn't a case compiles and only fails when the branch runs. Maybe use that wording here: a choice that isn't a task of this Dag doesn't compile, and one that isn't a case fails at run time. ########## java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt: ########## @@ -239,6 +243,59 @@ class TaskTest { Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs) } + @Test + @DisplayName("Should push the case a branch chose and skip every other one") + fun shouldSkipTheCasesNotChosen() { + TestBranch.choice = HandleLongTask::class.java + val transport = RecordingTransport() + + val result = runTask(branchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport)) + + Assertions.assertInstanceOf(SucceedTask::class.java, result) + Assertions.assertEquals(listOf("handle_short"), transport.skipped) + Assertions.assertEquals( + listOf("skipmixin_key" to mapOf("skipped" to listOf("handle_short")), "return_value" to "handle_long"), + transport.xComs, + ) + } + + @Test + @DisplayName("Should fail a branch that chose a task it cannot run") + fun shouldFailBranchThatChoseAnUnknownCase() { + TestBranch.choice = SuccessTask::class.java + val transport = RecordingTransport() + + val result = runTask(branchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport)) + + Assertions.assertInstanceOf(TaskState::class.java, result) + Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state) + Assertions.assertEquals(emptyList<String>(), transport.skipped) + Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs) + } + + @Test + @DisplayName("Should let a generated branch name its case by task ID") + fun shouldChooseByTaskId() { Review Comment: Could this get a failing twin like `shouldFailBranchThatChoseAnUnknownCase`? Nothing drives `byId` to its "not one of its cases" throw, and that throw is the only check on an annotation-style choice, since every task of the Dag has a `TaskIds` constant. ########## java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt: ########## @@ -0,0 +1,175 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk + +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertThrows +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +/** A case of the branches under test; a second class so a branch can tell two cases apart. */ +internal class HandleLong : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +/** A branch that chooses whatever [choice] holds. */ +class PickPath : BranchTask { + override fun choose( + context: Context, + client: Client, + ): Class<out Task> = choice + + companion object { + var choice: Class<out Task> = HandleLong::class.java + } +} + +internal class BranchTest { + private fun dagWithCases(): Triple<DagDef, BranchRef, Pair<TaskRef<Unit>, TaskRef<Unit>>> { + val dag = DagDef("d") + val handleLong = dag.task<Unit>("handle_long", HandleLong::class.java) + val handleShort = dag.task<Unit>("handle_short", NoopTask::class.java) + return Triple(dag, dag.Branch(PickPath::class.java), handleLong to handleShort) + } + + @Test + @DisplayName("Should take the branch's task ID from its class when none is given") + fun shouldDeriveTaskIdFromClass() { + val (dag, branch, cases) = dagWithCases() + branch.option(cases.first) + + assertEquals("pickPath", branch.id) + assertEquals(setOf(dag.tasks.getValue("pickPath")), cases.first.def.upstreams) + } + + @Test + @DisplayName("Should reject naming the same case twice") + fun shouldRejectDuplicateCase() { + val (_, branch, cases) = dagWithCases() + branch.option(cases.first) + + val error = assertThrows(IllegalArgumentException::class.java) { branch.option(cases.first) } + + assertEquals( + "Branch 'pickPath' already chooses between 'handle_long' and others; name each case once", + error.message, + ) + } + + @Test + @DisplayName("Should reject two cases a class-named branch could not tell apart") + fun shouldRejectCasesSharingAClass() { + val dag = DagDef("d") + val first = dag.task<Unit>("first", HandleLong::class.java) + val second = dag.task<Unit>("second", HandleLong::class.java) + val branch = dag.Branch(PickPath::class.java).option(first) + + val error = assertThrows(IllegalArgumentException::class.java) { branch.option(second) } + + assertEquals( + "Branch 'pickPath' cannot choose between 'first' and 'second': both run " + + "'org.apache.airflow.sdk.HandleLong', and a branch names its case by class", + error.message, + ) + } + + @Test + @DisplayName("Should reject a case declared in another Dag") + fun shouldRejectCaseFromAnotherDag() { + val (_, branch, _) = dagWithCases() + val other = DagDef("other").task<Unit>("handle_long", HandleLong::class.java) + + val error = assertThrows(IllegalArgumentException::class.java) { branch.option(other) } + + assertEquals( + "Branch 'pickPath' of Dag 'd' cannot run task 'handle_long' of Dag 'other'; " + + "name a task of the same Dag", + error.message, + ) + } + + @Test + @DisplayName("Should reject registering a branch with nothing to choose between") + fun shouldRejectBranchWithoutCases() { + val (dag, _, _) = dagWithCases() + + val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) } + + assertEquals( + "Branch 'pickPath' has no task to choose between; call option(...), in Dag 'd'", + error.message, + ) + } + + @Test + @DisplayName("Should reject a branch over a task that chooses nothing") + fun shouldRejectBranchOverAPlainTask() { + val dag = DagDef("d") + + val error = + assertThrows(IllegalArgumentException::class.java) { + BranchRef.of(dag.task<Unit>("plain", NoopTask::class.java)) + } + + assertEquals( + "Task 'plain' runs 'org.apache.airflow.sdk.NoopTask', which chooses nothing; " + + "a branch runs a BranchTask", + error.message, + ) + } + + @Test + @DisplayName("Should reject naming a case after the Dag was registered") + fun shouldRejectCaseAfterRegistration() { + val (dag, branch, cases) = dagWithCases() + branch.option(cases.first) + Bundle().register(dag) + + val error = assertThrows(IllegalArgumentException::class.java) { branch.option(cases.second) } + + assertEquals( + "Branch 'pickPath' of Dag 'd' is already registered; name every case before the Dag is " + + "added to a Bundle", + error.message, + ) + } + + @Test + @DisplayName("Should declare a branch inside a task group under the group's prefix") + fun shouldDeclareBranchInGroup() { + val dag = DagDef("d") + val reports = dag.taskGroup("reports") + val long = reports.task<Unit>("long", HandleLong::class.java) + val branch = reports.Branch("pick", PickPath::class.java).option(long) + + assertEquals("reports.pick", branch.id) + assertEquals(listOf("reports.long", "reports.pick"), dag.tasks.keys.toList()) + } + + @Test + @DisplayName("Should carry the task ID of a case rather than its class") + fun taskIdNamesOneTask() { Review Comment: This only checks `TaskId`'s own `equals` and `toString`, so it doesn't test what the name says. In the same file, `PickPath.choice` is a mutable global nothing here assigns; `choose` could just return `HandleLong::class.java`. ########## java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt: ########## @@ -0,0 +1,291 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk + +import kotlin.Throws + +/** + * The Airflow task ID of one task of a Dag. + * + * A `@Builder.Branch` method names the task it chose with one of these. The + * generated `<Dag>Builder.TaskIds` holds a constant per task of the Dag, so + * the choice is checked where it is written rather than when the task runs. + * + * @property value The task ID. + */ +class TaskId private constructor( + val value: String, +) { + companion object { + /** + * @suppress + * + * The ID [value] names. Public so a generated `TaskIds` holder can declare + * its constants; user code names a task through those constants. + */ + @JvmStatic + fun of(value: String): TaskId = TaskId(value) + } + + override fun equals(other: Any?): Boolean = other is TaskId && other.value == value + + override fun hashCode(): Int = value.hashCode() + + override fun toString(): String = value +} + +/** + * A task that chooses one of several tasks to run; every other one is skipped. + * + * Register one with [DagDef.Branch], then list what it can choose with + * [BranchRef.option]. The choice is the task's own class, so javac checks it: + * + * ```java + * public class PickPath implements BranchTask { + * @Override + * public Class<? extends Task> choose(Context context, Client client) { + * return ((Number) client.getXCom("extract")).longValue() > 1000 + * ? HandleLong.class + * : HandleShort.class; + * } + * } + * ``` + * + * The SDK runs [choose] and pushes the chosen task's ID as this task's return + * value, so [execute] is never called. + * + * Because a branch names a case by its class, no two of its cases may be + * registered from the same class; registering the Dag reports that. Review Comment: Nit: this throws from the second `option(...)` call, not when the Dag is registered (the comment in `BranchDef.add` says so). -- 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]
