@@ -96,6 +96,7 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName
9696@SupportedAnnotationTypes(
9797 " org.apache.airflow.sdk.Builder.Dag" ,
9898 " org.apache.airflow.sdk.Builder.Task" ,
99+ " org.apache.airflow.sdk.Builder.If" ,
99100 " org.apache.airflow.sdk.Builder.TaskGroup" ,
100101 " org.apache.airflow.sdk.Builder.TaskHandler" ,
101102 " org.apache.airflow.sdk.Builder.Deps" ,
@@ -356,29 +357,42 @@ class BuilderProcessor : AbstractProcessor() {
356357 builderName : ClassName ,
357358 inGroup : Boolean ,
358359 ): MethodSpec {
360+ val condition = decl.kind == TaskKind .CONDITION
359361 val method =
360362 MethodSpec
361363 .methodBuilder(decl.method.simpleName.toString())
362364 .addModifiers(Modifier .PUBLIC , Modifier .DEFAULT )
363- .returns(ParameterizedTypeName .get(TASK_HANDLE_TYPE , TypeName .get(decl.method.returnType).boxIfPossible()))
365+ .returns(
366+ if (condition) {
367+ CONDITION_REF_TYPE
368+ } else {
369+ ParameterizedTypeName .get(TASK_HANDLE_TYPE , TypeName .get(decl.method.returnType).boxIfPossible())
370+ },
371+ )
364372 decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) }
365373 val def = taskDefCode(decl, CodeBlock .of($$" $T .$L " , builderName, decl.className))
366374 // The view knows the group it belongs to, so the recorder is told where
367375 // the task goes instead of deriving it from the task's ID.
368376 val group = if (inGroup) CodeBlock .of(" groupId()" ) else CodeBlock .of($$" $S " , " " )
377+ val wrap = { call: CodeBlock -> if (condition) CodeBlock .of($$" $T .of($L )" , CONDITION_REF_TYPE , call) else call }
369378 if (decl.dataParams.isEmpty()) {
370- method.addStatement($$" return $T .node($L , $L )" , REFS_TYPE , group, def)
379+ method.addStatement($$" return $L " , wrap( CodeBlock .of($$ " $ T .node($L , $L )" , REFS_TYPE , group, def)) )
371380 } else {
372381 // The parameter names ride along so the serialized Dag can name each
373382 // argument, as the binding spec ADR-0007 defines requires.
374383 method.addStatement(
375- $$" return $T .call($L , $L , $T .of($L ), $L )" ,
376- REFS_TYPE ,
377- group,
378- def,
379- LIST_TYPE ,
380- CodeBlock .join(decl.dataParams.map { CodeBlock .of($$" $S " , it.name) }, " , " ),
381- decl.dataParams.joinToString { it.name },
384+ $$" return $L " ,
385+ wrap(
386+ CodeBlock .of(
387+ $$" $T .call($L , $L , $T .of($L ), $L )" ,
388+ REFS_TYPE ,
389+ group,
390+ def,
391+ LIST_TYPE ,
392+ CodeBlock .join(decl.dataParams.map { CodeBlock .of($$" $S " , it.name) }, " , " ),
393+ decl.dataParams.joinToString { it.name },
394+ ),
395+ ),
382396 )
383397 }
384398 return method.build()
@@ -396,7 +410,7 @@ class BuilderProcessor : AbstractProcessor() {
396410 CodeBlock
397411 .builder()
398412 .add($$" new $T ($S , $L .class)" , TASK_DEF_TYPE , decl.id, classRef)
399- explicitConfig(decl.method, TASK_ANNOTATION , TASK_STRUCTURAL_ATTRIBUTES , SchemaFields .TASK ).forEach { (key, value) ->
413+ explicitConfig(decl.method, decl.kind.annotation , TASK_STRUCTURAL_ATTRIBUTES , SchemaFields .TASK ).forEach { (key, value) ->
400414 taskDef.add($$" .config($S , $L )" , key, value)
401415 }
402416 return taskDef.build()
@@ -421,16 +435,24 @@ class BuilderProcessor : AbstractProcessor() {
421435 val tasks = mutableListOf<TaskDeclaration >()
422436 for (inner in el.enclosedElements) {
423437 if (inner !is ExecutableElement ) continue
424- val ann = inner.getAnnotation(Builder .Task ::class .java) ? : continue
438+ val task = inner.getAnnotation(Builder .Task ::class .java)
439+ val condition = inner.getAnnotation(Builder .If ::class .java)
440+ if (task == null && condition == null ) continue
441+ require(task == null || condition == null ) {
442+ " Method '${inner.simpleName} ' carries both @Builder.Task and @Builder.If; a condition is a " +
443+ " task, so it carries @Builder.If alone"
444+ }
445+ val kind = if (condition != null ) TaskKind .CONDITION else TaskKind .TASK
425446 if (inner.isVarArgs) throw IllegalArgumentException (" Cannot create task from vararg function ${inner.simpleName} " )
426- val localId = ann.id.ifBlank { inner.simpleName.toString() }
447+ checkDeciderReturn(kind, inner)
448+ val localId = (condition?.id ? : task!! .id).ifBlank { inner.simpleName.toString() }
427449 require(tasks.none { it.method.simpleName.contentEquals(inner.simpleName) }) {
428450 " Class ${el.simpleName} overloads task method '${inner.simpleName} '; a method's name is the " +
429451 " name of its generated task class and of its wiring-view method, so rename one and keep its " +
430- " task id with @Builder.Task (id = \" $localId \" )"
452+ " task id with ${kind.spelling} (id = \" $localId \" )"
431453 }
432454 tasks + =
433- TaskDeclaration (inner, (path + localId).joinToString(" ." ), collectDataParams(inner), el, classPath)
455+ TaskDeclaration (inner, (path + localId).joinToString(" ." ), collectDataParams(inner), el, classPath, kind )
434456 }
435457
436458 val groups = mutableListOf<GroupDeclaration >()
@@ -571,7 +593,7 @@ class BuilderProcessor : AbstractProcessor() {
571593 val name = decl.method.simpleName.toString()
572594 require(name !in reserved) {
573595 " Task method '$name ' clashes with a member of the wiring view; rename the method and keep " +
574- " the task id with @Builder.Task (id = \" ${decl.id.substringAfterLast(' .' )} \" )"
596+ " the task id with ${decl.kind.spelling} (id = \" ${decl.id.substringAfterLast(' .' )} \" )"
575597 }
576598 }
577599 val accessors = mutableMapOf<String , GroupDeclaration >()
@@ -682,13 +704,31 @@ class BuilderProcessor : AbstractProcessor() {
682704 }
683705 }
684706
707+ /* *
708+ * Checks what a decider hands back, which is what the SDK reads its decision
709+ * from. A plain task may return anything, including nothing.
710+ */
711+ private fun checkDeciderReturn (
712+ kind : TaskKind ,
713+ method : ExecutableElement ,
714+ ) {
715+ if (kind != TaskKind .CONDITION ) return
716+ val returns = method.returnType
717+ val boolean = returns.kind == TypeKind .BOOLEAN || with (processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE ) }
718+ require(boolean) {
719+ " @Builder.If method '${method.simpleName} ' returns $returns , but a condition returns boolean: " +
720+ " true runs the task named by Then, false the one named by Else"
721+ }
722+ }
723+
685724 private fun buildTask (decl : TaskDeclaration ): TypeSpec {
725+ val condition = decl.kind == TaskKind .CONDITION
686726 val executeSpec =
687727 MethodSpec
688- .methodBuilder(" execute" )
728+ .methodBuilder(if (condition) " decide " else " execute" )
689729 .addAnnotation(Override ::class .java)
690730 .addModifiers(Modifier .PUBLIC )
691- .returns(TypeName .VOID )
731+ .returns(if (condition) TypeName . BOOLEAN else TypeName .VOID )
692732 .addParameter(CONTEXT_TYPE , " context" )
693733 .addParameter(CLIENT_TYPE , " client" )
694734 .addException(Exception ::class .java)
@@ -734,10 +774,11 @@ class BuilderProcessor : AbstractProcessor() {
734774 }
735775 }
736776
737- if (inner.returnType.kind == TypeKind .VOID ) {
738- $$" new $T ().$L ($L )"
739- } else {
740- $$" client.setXCom(new $T ().$L ($L ))"
777+ when {
778+ // The SDK pushes a condition's result itself, after it has skipped the side not taken.
779+ condition -> $$" return new $T ().$L ($L )"
780+ inner.returnType.kind == TypeKind .VOID -> $$" new $T ().$L ($L )"
781+ else -> $$" client.setXCom(new $T ().$L ($L ))"
741782 }.also {
742783 executeSpec.addStatement(
743784 it,
@@ -749,7 +790,7 @@ class BuilderProcessor : AbstractProcessor() {
749790
750791 return TypeSpec
751792 .classBuilder(decl.className)
752- .addSuperinterface(Task ::class .java)
793+ .addSuperinterface(if (condition) CONDITION_TASK_TYPE else ClassName .get( Task ::class .java) )
753794 .addModifiers(Modifier .PUBLIC , Modifier .FINAL , Modifier .STATIC )
754795 .addMethod(executeSpec.build())
755796 .build()
@@ -893,6 +934,7 @@ private class TaskDeclaration(
893934 val dataParams : List <DataParam >,
894935 val owner : TypeElement ,
895936 val classPath : List <String > = emptyList(),
937+ val kind : TaskKind = TaskKind .TASK ,
896938) {
897939 val className: String =
898940 (classPath + method.simpleName.toString().replaceFirstChar(Char ::uppercase)).joinToString(" _" )
@@ -924,12 +966,27 @@ private val DAG_SOURCE_TYPE = ClassName.get(DagSource::class.java)
924966private val REFS_TYPE = ClassName .get(Refs ::class .java)
925967private val ARG_TYPE = ClassName .get(Arg ::class .java)
926968private val TASK_HANDLE_TYPE = ClassName .get(TaskRef ::class .java)
969+ private val CONDITION_TASK_TYPE = ClassName .get(ConditionTask ::class .java)
970+ private val CONDITION_REF_TYPE = ClassName .get(ConditionRef ::class .java)
971+ private val BOXED_BOOLEAN_TYPE = ClassName .get(" java.lang" , " Boolean" )
927972private val DEPS_TYPE = ClassName .get(Deps ::class .java)
928973private val GROUP_TYPE = DEPS_TYPE .nestedClass(" TaskGroup" )
929974private val LIST_TYPE = ClassName .get(List ::class .java)
930975
931976private const val DAG_ANNOTATION = " org.apache.airflow.sdk.Builder.Dag"
932- private const val TASK_ANNOTATION = " org.apache.airflow.sdk.Builder.Task"
977+
978+ /* *
979+ * What a declared task is, and the annotation that declares it. A decider
980+ * differs from a plain task in what it implements and in what its wiring-view
981+ * method hands back, but carries the same configuration attributes.
982+ */
983+ private enum class TaskKind (
984+ val annotation : String ,
985+ val spelling : String ,
986+ ) {
987+ TASK (" org.apache.airflow.sdk.Builder.Task" , " @Builder.Task" ),
988+ CONDITION (" org.apache.airflow.sdk.Builder.If" , " @Builder.If" ),
989+ }
933990
934991private val RESERVED_VIEW_NAMES =
935992 setOf (
0 commit comments