Java SDK: Resolve a task's arguments from the Dag's own wiring

A Dag authored in Java has no Python call site, so the supervisor sends no
argument bindings for it and a task's data parameters had nothing to
resolve against. Where the Dag wired inputs for a task, those stand in;
a stub-backed task keeps reading its bindings, including when the call
site bound none.

The authoring surface that records those inputs lands in the next commit,
so nothing declares them yet.

Author: jason810496

java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt

@@ -361,7 +361,7 @@ class BuilderProcessor : AbstractProcessor() {
       val paramType = TypeName.get(param.type)
       if (param.isTaskInput) {
         executeSpec.addStatement(
-          $$"$T $L = $T.bindInput(client, $T.class)",
+          $$"$T $L = $T.bindInput(context, client, $T.class)",
           paramType,
           param.local,
           ARG_VALUES_TYPE,

java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt

@@ -480,7 +480,7 @@ class BuilderTest {
            public static final class Named implements Task {
              @Override
              public void execute(Context context, Client client) throws Exception {
-               TestExample.ScoreInput context_ = ArgValues.bindInput(client, TestExample.ScoreInput.class);
+               TestExample.ScoreInput context_ = ArgValues.bindInput(context, client, TestExample.ScoreInput.class);
                new TestExample().named(context_);
              }
            }
@@ -873,7 +873,7 @@ class BuilderTest {
            public static final class Score implements Task {
              @Override
              public void execute(Context context, Client client) throws Exception {
-               TestExample.ScoreInput input = ArgValues.bindInput(client, TestExample.ScoreInput.class);
+               TestExample.ScoreInput input = ArgValues.bindInput(context, client, TestExample.ScoreInput.class);
                client.setXCom(new TestExample().score(client, input));
              }
            }

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Arg.kt

@@ -20,13 +20,20 @@
 package org.apache.airflow.sdk
 
 /**
- * A value a task can be given, of which [TaskRef] — the output of an upstream
- * task — is the only form so far.
+ * A value a task can be given: the output of an upstream task, carried by the
+ * [TaskRef] that task's registration returned, or an inline constant.
+ *
+ * A constant has to be wrapped because a bare `Double` cannot implement this
+ * type: boxed types only, no primitives.
  *
  * @param T Type of the value.
  */
 sealed class Arg<T>
 
+internal class LiteralArg<T>(
+  internal val value: T?,
+) : Arg<T>()
+
 /**
  * The output of a registered task, and the task's place in the flow.
  *

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt

@@ -124,6 +124,9 @@ data class Context(
   @JvmField val dagRun: DagRun,
   @JvmField val ti: TaskInstance,
 ) {
+  /** Registration of the executing task; resolves wired data inputs. */
+  internal var taskDef: TaskDef? = null
+
   internal companion object {
     fun from(request: StartupDetails) =
       Context(

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt

@@ -171,6 +171,7 @@ class TaskDef(
   }
 
   internal val configValues = linkedMapOf<String, Any>()
+  internal val inputs = mutableListOf<Arg<*>>()
   internal val upstreams = linkedSetOf<TaskDef>()
   internal var owner: DagDef? = null
 

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/InputTask.kt

@@ -57,7 +57,7 @@ interface InputTask<I : TaskInput> : Task {
   override fun execute(
     context: Context,
     client: Client,
-  ) = execute(context, client, ArgValues.bindInput(client, inputType()))
+  ) = execute(context, client, ArgValues.bindInput(context, client, inputType()))
 
   /**
    * Executes this task.

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt

@@ -74,9 +74,10 @@ internal object TaskRunner {
     request: StartupDetails,
     client: Client,
   ): Any {
-    val definition =
-      bundle.taskDef(request.ti.dagId, request.ti.taskId)?.definition
+    val taskDef =
+      bundle.taskDef(request.ti.dagId, request.ti.taskId)
         ?: return TaskResult.of(TaskState.State.REMOVED)
+    val definition = taskDef.definition
     val instance =
       try {
         definition.getDeclaredConstructor().newInstance()
@@ -103,7 +104,7 @@ internal object TaskRunner {
         return TaskResult.failure(request.tiContext.shouldRetry)
       }
     return try {
-      instance.execute(Context.from(request), client)
+      instance.execute(Context.from(request).also { it.taskDef = taskDef }, client)
       TaskResult.success()
     } catch (e: CancellationException) {
       throw e // Let coroutine cancellation propagate so the task coroutine unwinds.

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt

@@ -23,9 +23,13 @@ package org.apache.airflow.sdk.internal
 
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.databind.json.JsonMapper
+import org.apache.airflow.sdk.Arg
 import org.apache.airflow.sdk.Client
+import org.apache.airflow.sdk.Context
+import org.apache.airflow.sdk.LiteralArg
 import org.apache.airflow.sdk.MissingXComException
 import org.apache.airflow.sdk.TaskInput
+import org.apache.airflow.sdk.TaskRef
 import org.apache.airflow.sdk.execution.ArgBinding
 import org.apache.airflow.sdk.execution.Logger
 import java.lang.reflect.Field
@@ -42,6 +46,11 @@ import java.lang.reflect.Type
  * graph the scheduler ordered the run by. Flat data parameters resolve the
  * binding at their position (through [TaskArgs]); [TaskInput] fields resolve
  * bindings by name.
+ *
+ * A natively authored Dag has no stub call site, so the supervisor sends no
+ * bindings for it and the inputs the Dag itself wired stand in. When the
+ * supervisor sends bindings they are used for every parameter; the Dag's own
+ * inputs are read only when it sends none.
  */
 object ArgValues {
   private val mapper: ObjectMapper = JsonMapper.builder().build().findAndRegisterModules()
@@ -68,9 +77,17 @@ object ArgValues {
    */
   @JvmStatic
   fun <I : TaskInput> bindInput(
+    context: Context,
     client: Client,
     type: Class<I>,
   ): I {
+    // Runtime bindings carry argument names to match fields against. A wired
+    // input carries none, so it decodes into the whole input at once -- which
+    // is well defined because a TaskInput is a task's only data parameter.
+    wiredInputs(context, client)?.let { wired ->
+      return decode(resolveWired(wired[0], client), type) as I?
+        ?: throw missingInput(wired[0], type.simpleName)
+    }
     val input = newInput(type)
     val arguments = ArgIndex(client.argBindings)
     val unfilled = mutableListOf<String>()
@@ -147,6 +164,66 @@ object ArgValues {
     type: Type,
   ): Any? = decode(client.resolveBinding(binding), type)
 
+  /**
+   * Resolves the wired input at a data parameter's position into [type],
+   * passing null through.
+   */
+  internal fun valueWired(
+    input: Arg<*>,
+    client: Client,
+    type: Type,
+  ): Any? = decode(resolveWired(input, client), 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)
+      else ->
+        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.",
+        )
+      else -> MissingXComException("Input '$target' is wired to a null literal, so there is nothing to bind.")
+    }
+
+  private fun resolveWired(
+    input: Arg<*>,
+    client: Client,
+  ): Any? =
+    when (input) {
+      is TaskRef<*> -> client.getXCom(taskId = input.def.id)
+      is LiteralArg<*> -> input.value
+      else -> null
+    }
+
   /**
    * Builds the failure for a binding that resolved to nothing where a value is
    * required, naming [target] — the stub argument, or the [TaskInput] field

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/TaskArgs.kt

@@ -19,10 +19,12 @@
 
 package org.apache.airflow.sdk.internal
 
+import org.apache.airflow.sdk.Arg
 import org.apache.airflow.sdk.Client
 import org.apache.airflow.sdk.Context
 import org.apache.airflow.sdk.MissingXComException
 import org.apache.airflow.sdk.execution.ArgBinding
+import java.lang.reflect.Type
 
 /**
  * @suppress
@@ -46,6 +48,7 @@ class TaskArgs private constructor(
   private val context: Context,
   private val client: Client,
   private val arguments: List<ArgBinding>,
+  private val wired: List<Arg<*>>?,
 ) {
   companion object {
     /**
@@ -71,6 +74,13 @@ class TaskArgs private constructor(
       client: Client,
       declared: Int,
     ): TaskArgs {
+      ArgValues.wiredInputs(context, client)?.let { wired ->
+        check(wired.size == declared) {
+          "Task '${context.ti.taskId}' declares $declared data parameter(s) " +
+            "but the Dag wired ${wired.size} argument(s)"
+        }
+        return TaskArgs(context, client, emptyList(), wired)
+      }
       val bound = client.argBindings
       val arguments = if (bound.size == declared) bound else bound.filterNot { it.fromDefault }
       check(arguments.size == declared) {
@@ -83,7 +93,7 @@ class TaskArgs private constructor(
         "Task '${context.ti.taskId}' declares $declared data parameter(s) " +
           "but the stub call bound ${bound.size} argument(s)$defaults"
       }
-      return TaskArgs(context, client, arguments)
+      return TaskArgs(context, client, arguments, null)
     }
   }
 
@@ -96,7 +106,7 @@ class TaskArgs private constructor(
   fun <T : Any> get(
     position: Int,
     type: Class<T>,
-  ): T? = type.cast(ArgValues.valueAt(client, arguments[position], type))
+  ): T? = type.cast(valueAt(position, type))
 
   /**
    * Resolves the argument bound at [position] into the generic [type], passing
@@ -108,7 +118,7 @@ class TaskArgs private constructor(
   fun <T : Any> get(
     position: Int,
     type: TypeRef<T>,
-  ): T? = ArgValues.valueAt(client, arguments[position], type.type) as T?
+  ): T? = valueAt(position, type.type) as T?
 
   /**
    * Resolves the argument bound at [position] into [type], which must not be
@@ -136,7 +146,23 @@ class TaskArgs private constructor(
     type: TypeRef<T>,
   ): T = get(position, type) ?: throw missingAt(position)
 
-  // The stub signature's own parameter name is the clearest label for a failure
-  // here: it is what the Dag author has to change.
-  private fun missingAt(at: Int) = ArgValues.missing(arguments[at], context.ti.taskId)
+  private fun valueAt(
+    position: Int,
+    type: Type,
+  ): Any? =
+    if (wired != null) {
+      ArgValues.valueWired(wired[position], client, type)
+    } else {
+      ArgValues.valueAt(client, arguments[position], type)
+    }
+
+  // A bound argument is labelled with the stub signature's own parameter name,
+  // which is what the Dag author has to change. A wired one has no such name,
+  // so it is labelled with the position it feeds.
+  private fun missingAt(at: Int): MissingXComException =
+    if (wired != null) {
+      ArgValues.missingWired(wired[at], "position $at")
+    } else {
+      ArgValues.missing(arguments[at], context.ti.taskId)
+    }
 }

java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt

@@ -94,3 +94,18 @@ internal fun taskContext(): Context =
     dagRun = DagRun("d", "r", null, null, null, null, null, emptyMap()),
     ti = TaskInstance("d", "r", "t", null, 1),
   )
+
+internal class NoopTask : Task {
+  override fun execute(
+    context: Context,
+    client: Client,
+  ) = Unit
+}
+
+/** A context whose task was wired by its Dag with the given inputs. */
+internal fun contextWiredWith(inputs: List<Arg<*>>): Context {
+  val def = TaskDef("t", NoopTask::class.java)
+  DagDef("d").addTask(def)
+  def.inputs += inputs
+  return taskContext().also { it.taskDef = def }
+}