Java SDK: answer a task handler parse with every handler

The parse request no longer names Dags, and the Dag processor reuses an artifact's answer for every Dag file that uses the artifact, so the answer must not depend on what was asked. The vendored schema is resynced in the same commit because the request model generated from it no longer has dag_ids.

Author: jason810496

airflow-core/tests/unit/dag_processing/test_task_handler_processor_java.py

@@ -162,13 +162,6 @@ def test_a_built_bundle_declares_its_task_handlers(example_bundle):
         bundle_path=example_bundle,
         bundle_name="java-task-handlers",
         artifact_rel_path=jar.name,
-        dag_ids=[
-            "java_xcom_casting_example",
-            "java_interface_example",
-            "java_annotation_example",
-            "java_native_interface_example",
-            "not_registered",
-        ],
         logger=structlog.get_logger(),
     )
 
@@ -207,9 +200,9 @@ def test_a_built_bundle_declares_its_task_handlers(example_bundle):
         },
     )
     assert list(result.task_handlers) == [
-        "java_xcom_casting_example",
         "java_interface_example",
         "java_annotation_example",
+        "java_xcom_casting_example",
     ]
     with zipfile.ZipFile(jar) as zf:
         manifest = _parse_manifest(zf.read("META-INF/MANIFEST.MF"))

java-sdk/README.md

@@ -677,8 +677,7 @@ launches the JVM application as a subprocess. The flow is:
 The Dag processor also launches each JAR to ask which task handlers it
 registers. The first frame is then a `TaskHandlerParseRequest`, and
 `execution/TaskHandlerParse.kt` replies with one `TaskHandlerParsingResult`:
-the task-handler registrations of the requested Dags, each with its binding
-mode and parameters. Flat data parameters bind `positional` and a `TaskInput`
+every task-handler registration, each with its binding mode and parameters. Flat data parameters bind `positional` and a `TaskInput`
 binds `named`. The annotation processor records a handler's parameter names on
 its generated task class (`internal/TaskParams.kt`), because Java does not keep
 them at run time. The JVM waits for the reply's acknowledgement, then exits.

java-sdk/sdk/schema/schema.json

@@ -4199,19 +4199,12 @@
       "type": "object"
     },
     "TaskHandlerParseRequest": {
-      "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime which task handlers an artifact registers for the given Dags.",
+      "description": "Request for Task Handler Parsing.\n\nAsks a Lang-SDK runtime for every task handler an artifact registers.",
       "properties": {
         "file": {
           "title": "File",
           "type": "string"
         },
-        "dag_ids": {
-          "items": {
-            "type": "string"
-          },
-          "title": "Dag Ids",
-          "type": "array"
-        },
         "bundle_path": {
           "format": "path",
           "title": "Bundle Path",
@@ -4230,15 +4223,14 @@
       },
       "required": [
         "file",
-        "dag_ids",
         "bundle_path",
         "bundle_name"
       ],
       "title": "TaskHandlerParseRequest",
       "type": "object"
     },
     "TaskHandlerParsingResult": {
-      "description": "Result of Task Handler Parsing.\n\nThe task handlers a Lang-SDK artifact registers, keyed by Dag id.",
+      "description": "Result of Task Handler Parsing.\n\nEvery task handler a Lang-SDK artifact registers, keyed by Dag id.\n\nThe answer depends only on the artifact, never on the request.",
       "properties": {
         "fileloc": {
           "title": "Fileloc",

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

@@ -73,9 +73,8 @@ class ApiError(
  * execution requests to the registered [Bundle].
  *
  * The Dag processor starts the same entry point to ask which task handlers
- * the bundle registers. That process replies with the handlers registered for
- * the requested Dags and how each binds its arguments, runs no task, and
- * exits.
+ * the bundle registers. That process replies with every registered handler
+ * and how each binds its arguments, runs no task, and exits.
  *
  * The typical entry point is:
  *

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

@@ -36,29 +36,23 @@ import java.lang.reflect.Modifier
 import java.lang.reflect.Type
 
 /**
- * Answers a [TaskHandlerParseRequest] with the task handlers [bundle]
- * registers for the requested Dags, as a TaskHandlerParsingResult body.
+ * Answers a [TaskHandlerParseRequest] with every task handler [bundle]
+ * registers, keyed by Dag id in registration order, as a
+ * TaskHandlerParsingResult body.
  *
  * Only task-handler registrations count: a Dag declared in Java is its own,
- * not a Python Dag's. A requested Dag with no handlers is left out. The
- * generated models cannot hold `task_handlers`, so the body is a map.
+ * not a Python Dag's. The generated models cannot hold `task_handlers`, so
+ * the body is a map.
  */
 internal fun parseTaskHandlers(
   bundle: Bundle,
   request: TaskHandlerParseRequest,
-): Map<String, Any?> {
-  val declarations = linkedMapOf<String, List<Map<String, Any?>>>()
-  for (dagId in request.dagIds) {
-    if (dagId in declarations) continue
-    val dag = bundle.taskHandlers[dagId] ?: continue
-    declarations[dagId] = dag.tasks.values.map(::declare)
-  }
-  return linkedMapOf(
+): Map<String, Any?> =
+  linkedMapOf(
     "type" to "TaskHandlerParsingResult",
     "fileloc" to request.file,
-    "task_handlers" to declarations,
+    "task_handlers" to bundle.taskHandlers.mapValues { (_, dag) -> dag.tasks.values.map(::declare) },
   )
-}
 
 private fun declare(task: TaskDef): Map<String, Any?> {
   val (binding, params) = bindingOf(task.definition)

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

@@ -135,7 +135,7 @@ class ServerTest {
     val supervisor =
       Thread {
         runBlocking {
-          toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar", listOf("etl")))
+          toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar"))
           val prefix = fromServer.readByteArray(4)
           val raw = Frame.decodeRaw(ArrayBufferInput(fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt())))
           reported.put(raw)
@@ -169,7 +169,7 @@ class ServerTest {
     val supervisor =
       Thread {
         runBlocking {
-          toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar", listOf("etl")))
+          toServer.writeFrame(taskHandlerParseRequestFrame(3, "/bundle/etl.jar"))
           val prefix = fromServer.readByteArray(4)
           fromServer.readByteArray(Frame.parseLengthPrefix(prefix).toInt())
           toServer.close()
@@ -187,17 +187,14 @@ class ServerTest {
   private fun taskHandlerParseRequestFrame(
     id: Int,
     file: String,
-    dagIds: List<String>,
   ): ByteArray {
     val out = ByteArrayOutputStream()
     MessagePack.newDefaultPacker(out).use { packer ->
       packer.packArrayHeader(2)
       packer.packInt(id)
-      packer.packMapHeader(5)
+      packer.packMapHeader(4)
       packer.packString("type").packString("TaskHandlerParseRequest")
       packer.packString("file").packString(file)
-      packer.packString("dag_ids").packArrayHeader(dagIds.size)
-      dagIds.forEach { packer.packString(it) }
       packer.packString("bundle_path").packString("/bundle")
       packer.packString("bundle_name").packString("java-task-handlers")
     }

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

@@ -109,10 +109,9 @@ private class Audit : Task {
   ) = Unit
 }
 
-private fun request(vararg dagIds: String) =
+private fun request() =
   TaskHandlerParseRequest().apply {
     file = "/bundles/java/etl.jar"
-    this.dagIds = dagIds.toList()
     bundlePath = "/bundles/java"
     bundleName = "java-task-handlers"
   }
@@ -151,8 +150,8 @@ private fun assertConforms(
 
 class TaskHandlerParseTest {
   @Test
-  @DisplayName("Should declare each requested Dag's task handlers in registration order")
-  fun declaresRequestedDagsInRegistrationOrder() {
+  @DisplayName("Should declare every task handler in registration order")
+  fun declaresEveryTaskHandlerInRegistrationOrder() {
     val bundle =
       Bundle()
         .register("etl", "extract", GeneratedFlat::class.java)
@@ -163,7 +162,7 @@ class TaskHandlerParseTest {
         .register("report", "audit", Audit::class.java)
         .register(DagDef("native").addTask("audit", Audit::class.java))
 
-    val result = parseTaskHandlers(bundle, request("etl", "native", "missing", "etl"))
+    val result = parseTaskHandlers(bundle, request())
 
     val regions =
       mapOf("anyOf" to listOf(mapOf("type" to "array", "items" to NULLABLE_STRING), mapOf("type" to "null")))
@@ -189,9 +188,11 @@ class TaskHandlerParseTest {
                 mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList<Any>()),
                 mapOf("task_id" to "notify", "binding" to "named", "params" to emptyList<Any>()),
               ),
+            "report" to listOf(mapOf("task_id" to "audit", "binding" to "named", "params" to emptyList<Any>())),
           ),
       )
     assertEquals(expected, result)
+    assertEquals(listOf("etl", "report"), (result["task_handlers"] as Map<*, *>).keys.toList())
   }
 
   @Test
@@ -203,7 +204,7 @@ class TaskHandlerParseTest {
         .register("etl", "score", GeneratedInput::class.java)
         .register("etl", "audit", Audit::class.java)
 
-    val result = parseTaskHandlers(bundle, request("etl"))
+    val result = parseTaskHandlers(bundle, request())
 
     assertConforms("TaskHandlerParsingResult", result)
     val bindings =
@@ -221,11 +222,11 @@ class TaskHandlerParseTest {
   }
 
   @Test
-  @DisplayName("Should send an empty map when no requested Dag has task handlers")
-  fun sendsEmptyTaskHandlersWhenNothingMatches() {
-    val bundle = Bundle().register("etl", "audit", Audit::class.java)
+  @DisplayName("Should send an empty map when the bundle registers no task handler")
+  fun sendsEmptyTaskHandlersWithoutHandlers() {
+    val bundle = Bundle().register(DagDef("native").addTask("audit", Audit::class.java))
 
-    val result = parseTaskHandlers(bundle, request("other"))
+    val result = parseTaskHandlers(bundle, request())
 
     assertEquals(emptyMap<String, Any?>(), result["task_handlers"])
   }