guan404ming commented on code in PR #73596:
URL: https://github.com/apache/airflow/pull/73596#discussion_r4144907628
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt:
##########
@@ -147,6 +173,113 @@ object ArgValues {
type: Type,
): Any? = decode(client.resolveBinding(binding), type)
+ /**
+ * Reports a Dag that wired more or fewer inputs than a [TaskInput] can take.
+ * A [TaskInput] is a task's only data parameter, so exactly one input feeds
+ * it; a mismatch warns and the first input is used, leaving the task to run
+ * on what it can bind rather than failing the run outright.
+ */
+ private fun warnWiredArity(
+ client: Client,
+ type: Class<*>,
+ wired: Int,
+ ) {
+ if (wired == 1) return
+ logger.warning(
+ "Dag wired a different number of argument(s) than the task handler
declares",
+ mapOf(
+ "task_id" to client.details.ti.taskId,
+ "input" to type.simpleName,
+ "declared" to 1,
+ "wired" to wired,
+ ),
+ )
+ }
+
+ /**
+ * Resolves every wired input at once: each upstream is read once however
+ * many parameters it feeds, and several upstreams are read concurrently.
+ * The supervisor protocol matches responses to requests by id, so a task
+ * wired to several upstreams waits roughly one round trip rather than one
+ * per parameter.
+ */
+ internal fun resolveWiredAll(
+ inputs: List<Arg<*>>,
+ client: Client,
+ ): List<Any?> {
+ val upstreams = inputs.filterIsInstance<TaskRef<*>>().map { it.def.id
}.distinct()
+ val fetched =
+ when (upstreams.size) {
+ 0 -> emptyMap()
+ 1 -> mapOf(upstreams[0] to client.getXCom(taskId = upstreams[0]))
+ else ->
+ runBlocking {
+ upstreams.map { taskId -> async(Dispatchers.IO) { taskId to
client.getXCom(taskId = taskId) } }.awaitAll()
+ }.toMap()
+ }
+ return inputs.map { input ->
+ when (input) {
+ is TaskRef<*> -> fetched[input.def.id]
+ is LiteralArg<*> -> input.value
+ }
+ }
+ }
+
+ /** Decodes an already-resolved wired value into [type], passing null
through. */
+ internal fun decodeWired(
+ value: Any?,
+ type: Type,
+ ): Any? = decode(value, type)
+
+ /**
+ * The inputs the Dag wired for this task, or null when the run's arguments
+ * come from the stub call site. A task with no wired inputs reads the
+ * bindings, so a stub call that bound nothing keeps its own diagnostics.
+ */
+ internal fun wiredInputs(
+ context: Context,
+ client: Client,
+ ): List<Arg<*>>? = if (client.argBindings.isEmpty())
context.taskDef?.inputs?.takeIf { it.isNotEmpty() } else null
+
+ /**
+ * The failure for a wired argument that resolved to nothing where a value is
+ * required, naming [target] — the position of the parameter it feeds.
+ */
+ internal fun missingWired(
+ input: Arg<*>,
+ target: String,
+ ): MissingXComException =
+ when (input) {
+ is TaskRef<*> -> MissingXComException(input.def.id, target)
+ is LiteralArg<*> ->
+ MissingXComException(
+ "Task parameter '$target' is wired to a null literal, but has a
primitive type that cannot " +
+ "be null; declare a boxed type (e.g. Integer instead of int) to
receive null.",
+ )
+ }
+
+ /** The failure for a wired input that resolved to nothing for a
[TaskInput]. */
+ private fun missingInput(
+ input: Arg<*>,
+ target: String,
+ ): MissingXComException =
+ when (input) {
+ is TaskRef<*> ->
+ MissingXComException(
+ "Input '$target' requires an XCom from task '${input.def.id}', but
none was pushed.",
+ )
+ is LiteralArg<*> -> MissingXComException("Input '$target' is wired to a
null literal, so there is nothing to bind.")
+ }
+
+ private fun resolveWired(
Review Comment:
`resolveWired` matches `resolveWiredAll` for a single input, so we could
reuse that one instead.
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt:
##########
@@ -123,9 +123,14 @@ data class TaskInstance(
data class Context(
@JvmField val dagRun: DagRun,
@JvmField val ti: TaskInstance,
+ /** Registration of the executing task; resolves wired data inputs. */
+ internal val taskDef: TaskDef? = null,
Review Comment:
Adding `taskDef` here changes the public data class constructor, `equals`
and `toString`. Keeping it outside would avoid that.
--
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]
