jason810496 commented on code in PR #74439:
URL: https://github.com/apache/airflow/pull/74439#discussion_r4218527127
##########
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:
Resolved in 141e5cedb9c: the TaskIds holder and both of its name checks are
gone, so a task id starting with a digit no longer stops a Dag compiling.
##########
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:
Done in 141e5cedb9c: a case is named by its class in both styles now, so
TaskId, the generated TaskIds and the hidden id-based interface are gone along
with the string escape hatch; the trade is that javac only checks the choice is
a Task, which the KDoc now says.
##########
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:
Done in 2594c90237e: it is DagDef.Switch and SwitchRef.Case now, capitalised
for the same reason If is.
##########
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:
Done in 35fd6f24a51: the wording now says the second Case call throws.
##########
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:
Done in b10b0dc0595: the test asserts the decision carries the chosen task's
id and skips the other, and the decider no longer reads a global nothing
assigns.
##########
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:
Done in 08909559642, covering the generated shape choosing a task of the Dag
that is not one of its cases.
##########
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:
Done in 35fd6f24a51, using that wording.
--
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]
