henry3260 commented on code in PR #71189:
URL: https://github.com/apache/airflow/pull/71189#discussion_r4183415938
##########
java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt:
##########
@@ -190,46 +228,201 @@ class BuilderProcessor : AbstractProcessor() {
explicitConfig(el, DAG_ANNOTATION, DAG_STRUCTURAL_ATTRIBUTES,
SchemaFields.DAG).forEach { (key, value) ->
buildMethod.addStatement($$"dag.config($S, $L)", key, value)
}
+ buildMethod.addStatement(
+ $$"return $T.record(dag, $T.of($L), new $T()::depends)",
+ REFS_TYPE,
+ ClassName.get(List::class.java),
+ CodeBlock.join(declarations.map { CodeBlock.of($$"$S", it.id) }, ", "),
+ ClassName.get(deps),
+ )
+ builderClass.addMethod(buildMethod.build())
- for (inner in el.enclosedElements) {
- if (inner !is ExecutableElement) continue
- if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task
from vararg function ${inner.simpleName}")
-
- val taskAnn = inner.getAnnotation(Builder.Task::class.java) ?: continue
- val innerName =
inner.simpleName.toString().replaceFirstChar(Char::uppercase)
+ declarations.forEach { builderClass.addType(buildTask(it, el)) }
+ return builderClass.build()
+ }
- builderClass.addType(buildTask(innerName, inner, el))
+ /**
+ * Generates the Dag's wiring view: one default method per task, with the
+ * injected arguments stripped, each data argument lifted to [Arg] and the
+ * return lifted to [TaskRef].
+ *
+ * It is an interface so the `@Builder.Deps` class can *implement* it and
+ * keep its own `extends` free, and so the Dag class's real task methods --
+ * which differ only in their injected arguments -- do not clash with it.
+ */
+ private fun buildDeps(
+ el: TypeElement,
+ declarations: List<TaskDeclaration>,
+ builderName: ClassName,
+ depsName: ClassName,
+ ): TypeSpec {
+ val view =
+ TypeSpec
+ .interfaceBuilder(depsName)
+ .addModifiers(Modifier.PUBLIC)
+ .addSuperinterface(DEPS_TYPE)
+ .addJavadoc(
+ "Wiring view of {@link \$T}'s task methods, for declaring its task
graph.\n\n" +
+ "<p>Calling one registers its task with the Dag being built;
passing the handle it\n" +
+ "returned into another call feeds the upstream's output into that
task's parameter\n" +
+ "and wires the data edge. {@code before} and {@code after} wire an
ordering-only edge.\n",
+ ClassName.get(el),
+ )
- buildMethod.addStatement(
- $$"dag.addTask($L)",
- taskDefCode(inner, taskAnn.id.ifBlank { inner.simpleName.toString() },
innerName),
- )
+ for (decl in declarations) {
+ val method =
+ MethodSpec
+ .methodBuilder(decl.method.simpleName.toString())
+ .addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
+ .returns(ParameterizedTypeName.get(TASK_HANDLE_TYPE,
TypeName.get(decl.method.returnType).boxIfPossible()))
+ decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) }
+ val def = taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName,
decl.className))
+ if (decl.dataParams.isEmpty()) {
+ method.addStatement($$"return $T.node($L)", REFS_TYPE, def)
+ } else {
+ method.addStatement(
+ $$"return $T.call($L, $L)",
+ REFS_TYPE,
+ def,
+ decl.dataParams.joinToString { it.name },
+ )
+ }
+ view.addMethod(method.build())
}
-
- buildMethod.addStatement("return dag")
- builderClass.addMethod(buildMethod.build())
- return builderClass.build()
+ return view.build()
}
/**
- * Emits `new TaskDef(id, <className>.class)` with the explicitly-written
+ * Emits `new TaskDef(id, <classRef>.class)` with the explicitly-written
* `@Builder.Task` attributes lowered into chained `.config` calls.
*/
private fun taskDefCode(
- method: ExecutableElement,
- id: String,
- className: String,
+ decl: TaskDeclaration,
+ classRef: CodeBlock,
): CodeBlock {
val taskDef =
CodeBlock
.builder()
- .add($$"new $T($S, $L.class)", TASK_DEF_TYPE, id, className)
- explicitConfig(method, TASK_ANNOTATION, TASK_STRUCTURAL_ATTRIBUTES,
SchemaFields.TASK).forEach { (key, value) ->
+ .add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef)
+ explicitConfig(decl.method, TASK_ANNOTATION, TASK_STRUCTURAL_ATTRIBUTES,
SchemaFields.TASK).forEach { (key, value) ->
taskDef.add($$".config($S, $L)", key, value)
}
return taskDef.build()
}
+ /**
+ * Maps a data parameter's declared type to its wiring-view input type,
+ * `Arg<? extends T>` of the boxed type. A numeric parameter therefore takes
+ * only its own type, so javac rejects wiring that could lose a value, such
+ * as a `double` upstream into a `long` parameter. An `Object` parameter
+ * takes any upstream, including a `void` task's handle, whose value is null.
+ */
+ private fun inType(paramType: TypeMirror): TypeName =
+ ParameterizedTypeName.get(ARG_TYPE,
WildcardTypeName.subtypeOf(TypeName.get(paramType).boxIfPossible()))
+
+ private fun collectTasks(el: TypeElement): List<TaskDeclaration> {
+ val declarations = mutableListOf<TaskDeclaration>()
+ for (inner in el.enclosedElements) {
+ if (inner !is ExecutableElement) continue
+ val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue
+ if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task
from vararg function ${inner.simpleName}")
+ val id = ann.id.ifBlank { inner.simpleName.toString() }
+ require(declarations.none { it.id == id }) { "Tasks in Dag have
duplicate ID: $id" }
+ declarations += TaskDeclaration(inner, id, collectDataParams(inner))
+ }
+ return declarations
+ }
+
+ /**
+ * Finds and validates the class's `@Builder.Deps` wiring class, which
+ * declares the Dag's task graph and is what makes it a Dag Java owns.
+ *
+ * The generated builder runs `new Wiring()::depends`, so everything that
+ * expression needs is checked here, where the error can name the class.
+ */
+ private fun findDeps(
+ el: TypeElement,
+ view: ClassName,
+ ): TypeElement {
+ val classes =
+ el.enclosedElements
+ .filterIsInstance<TypeElement>()
+ .filter { it.getAnnotation(Builder.Deps::class.java) != null }
+ require(classes.isNotEmpty()) {
+ "Dag class ${el.simpleName} must declare a @Builder.Deps class
implementing ${view.simpleName()} " +
+ "to declare its task graph; a class of task bodies for a Dag the
Python file owns carries " +
+ "@Builder.TaskHandler instead"
+ }
+ val deps =
+ classes.singleOrNull()
+ ?: throw IllegalArgumentException(
+ "Dag class ${el.simpleName} declares more than one @Builder.Deps
class: " +
+ classes.joinToString { it.simpleName.toString() },
+ )
+ val name = deps.simpleName
+ require(deps.kind == ElementKind.CLASS && Modifier.ABSTRACT !in
deps.modifiers) {
+ "@Builder.Deps '$name' must be a concrete class"
+ }
+ require(Modifier.STATIC in deps.modifiers && Modifier.PRIVATE !in
deps.modifiers) {
+ "@Builder.Deps class '$name' must be static and non-private"
+ }
+ require(deps.interfaces.any { it.isView(view) }) {
+ "@Builder.Deps class '$name' must implement ${view.simpleName()}, the
wiring view of ${el.simpleName}"
+ }
+ require(
+ deps.enclosedElements
+ .filterIsInstance<ExecutableElement>()
+ .any { it.kind == ElementKind.CONSTRUCTOR && it.parameters.isEmpty()
&& Modifier.PRIVATE !in it.modifiers },
+ ) {
+ "@Builder.Deps class '$name' needs a non-private no-argument constructor"
+ }
+ val depends =
+ processingEnv.elementUtils
+ .getAllMembers(deps)
+ .filterIsInstance<ExecutableElement>()
+ .firstOrNull { it.isNoArgDepends() }
+ ?: throw IllegalArgumentException(
+ "@Builder.Deps class '$name' must have a non-private, no-argument
depends() method",
+ )
+ val checked = depends.thrownTypes.filterNot { isUnchecked(it) }
+ require(checked.isEmpty()) {
+ "depends() of @Builder.Deps class '$name' must not throw checked
exceptions: ${checked.joinToString()}"
+ }
+ return deps
+ }
+
+ /**
+ * Rejects a task method whose wiring-view twin would clash with a member
+ * the view or the wiring class already has: `depends`, `lit`, or a method
+ * of `Object`.
+ */
+ private fun checkViewName(decl: TaskDeclaration) {
+ val name = decl.method.simpleName.toString()
+ require(name !in RESERVED_VIEW_NAMES) {
+ "Task method '$name' clashes with a member of the wiring view; rename
the method and keep " +
+ "the task id with @Builder.Task(id = \"${decl.id}\")"
+ }
+ }
Review Comment:
Two task methods with the same name and one data parameter each generate
`extract(Arg)` twice in the view, and the duplicate is reported against the
generated file. Should `checkViewName` also check name plus parameter count?
--
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]
