Go SDK: answer TaskHandlerParseRequest with registered handlers

airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst

@@ -502,8 +502,9 @@ Deploying
 Copy or mount the packed bundle into the Dag bundle named by the coordinator's ``task_handler_bundle_name``.
 The :class:`~airflow.sdk.coordinators.executable.ExecutableCoordinator` scans that Dag bundle recursively,
 matches the incoming ``dag_id`` against each bundle's manifest, verifies the bundle's integrity hash, and
-launches the matching bundle. Bundles are identified by the trailer magic, not by filename (no extension on
-Linux/macOS, ``.exe`` on Windows), so the file name on the worker is irrelevant.
+launches the matching bundle. If multiple usable bundles declare the same Dag, the first in sorted path order
+wins. Bundles are identified by the trailer magic, not by filename (no extension on Linux/macOS, ``.exe`` on
+Windows), so the file name on the worker is irrelevant.
 
 Only files with the executable bit set are considered, so the packed bundles need a Dag bundle that keeps
 it. A ``LocalDagBundle`` over a directory you manage does; object-store Dag bundles such as ``S3DagBundle``

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

@@ -0,0 +1,152 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Pack the Go SDK example bundle and probe it for its task handlers."""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+from typing import TYPE_CHECKING
+from unittest import mock
+
+import pytest
+import structlog
+
+from airflow.dag_processing.processor import TaskHandlerDeclaration, TaskHandlerParam
+from airflow.dag_processing.task_handler_processor import LangSDKTaskHandlerProcessorProcess
+from airflow.sdk.coordinators._subprocess import supports_task_handler_parsing
+from airflow.sdk.execution_time import supervisor
+from airflow.sdk.execution_time.coordinator import get_coordinator_manager, reset_coordinator_manager
+
+from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.paths import AIRFLOW_ROOT_PATH
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("AIRFLOW_LANG_SDK_REAL_PROBE_TESTS") != "1",
+    reason="set AIRFLOW_LANG_SDK_REAL_PROBE_TESTS=1 to build and probe a real Go bundle",
+)
+
+STRING = {"type": "string"}
+INT64 = {"type": "integer", "format": "int64"}
+DOUBLE = {"type": "number", "format": "double"}
+
+
+def _make_nullable(schema: dict) -> dict:
+    return {"anyOf": [schema, {"type": "null"}]}
+
+
+@pytest.fixture(scope="module")
+def go_bundle(tmp_path_factory) -> Path:
+    if shutil.which("go") is None:
+        pytest.skip("needs a Go toolchain on PATH")
+    bundle = tmp_path_factory.mktemp("go-task-handlers") / "example_dags"
+    completed = subprocess.run(
+        ["go", "tool", "airflow-go-pack", "--output", os.fspath(bundle), "./example/bundle"],
+        cwd=AIRFLOW_ROOT_PATH / "go-sdk",
+        env={**os.environ, "CGO_ENABLED": "0"},
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+    assert completed.returncode == 0, completed.stderr
+    return bundle
+
+
+@pytest.fixture(autouse=True)
+def _go_coordinator():
+    spec = {"go": {"classpath": "airflow.sdk.coordinators.executable.ExecutableCoordinator"}}
+    reset_coordinator_manager()
+    try:
+        # A bare fork, so the parse child sees this config.
+        with (
+            conf_vars({("sdk", "coordinators"): json.dumps(spec)}),
+            mock.patch.object(supervisor, "_should_use_exec", autospec=True, return_value=False),
+        ):
+            yield
+    finally:
+        reset_coordinator_manager()
+
+
+def test_probes_the_task_handlers_of_a_packed_go_bundle(go_bundle):
+    coordinator = get_coordinator_manager().get_coordinator("go")
+    artifact = coordinator._find_task_handler_artifact(bundle_path=go_bundle.parent, dag_id="simple_dag")
+    assert supports_task_handler_parsing(artifact.schema_version)
+    assert artifact.path == go_bundle.resolve()
+
+    result = LangSDKTaskHandlerProcessorProcess.run(
+        coordinator="go",
+        path=artifact.path,
+        bundle_path=go_bundle.parent,
+        bundle_name="go-task-handlers",
+        artifact_rel_path=os.fspath(artifact.path.relative_to(go_bundle.parent.resolve())),
+        logger=structlog.get_logger(),
+    )
+
+    assert result.import_errors is None
+    assert set(result.task_handlers) == {
+        "simple_dag",
+        "concurrent_xcom_dag",
+        "taskflow_binding_dag",
+        "variable_write_dag",
+    }
+    assert [d.task_id for d in result.task_handlers["simple_dag"]] == ["extract", "transform", "load"]
+
+    declarations = {d.task_id: d for d in result.task_handlers["taskflow_binding_dag"]}
+    assert declarations["via_flat_args"] == TaskHandlerDeclaration(
+        task_id="via_flat_args",
+        binding="positional",
+        params=[
+            TaskHandlerParam(name=None, value_schema=schema)
+            for schema in (
+                STRING,
+                INT64,
+                DOUBLE,
+                {"type": "boolean"},
+                _make_nullable({"type": "array", "items": STRING}),
+                {"type": "object"},
+                _make_nullable({"type": "array", "items": INT64}),
+                _make_nullable(STRING),
+            )
+        ],
+    )
+    assert declarations["via_struct_arg_tag"] == TaskHandlerDeclaration(
+        task_id="via_struct_arg_tag",
+        binding="named",
+        params=[
+            TaskHandlerParam(name="region_code", exact_name=True, value_schema=STRING),
+            TaskHandlerParam(name="threshold", exact_name=True, value_schema=DOUBLE),
+        ],
+    )
+    # The Dag passes this handler an argument it does not take, which only warns under named binding.
+    assert declarations["via_struct_more_args"] == TaskHandlerDeclaration(
+        task_id="via_struct_more_args",
+        binding="named",
+        params=[TaskHandlerParam(name="region_code", exact_name=True, value_schema=STRING)],
+    )
+    assert declarations["via_flat_map"] == TaskHandlerDeclaration(
+        task_id="via_flat_map",
+        binding="named",
+        params=[
+            TaskHandlerParam(name="Region", value_schema=STRING),
+            TaskHandlerParam(name="Count", value_schema=INT64),
+        ],
+    )

contributing-docs/30_new_language_sdk.rst

@@ -126,6 +126,43 @@ or the task's own Dag bundle when it is unset, pinned for the whole task.
 Subclasses should scan those directories rather than locating artifacts
 themselves.
 
+SubprocessCoordinator: answering a task handler parse
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When a Python Dag file has stub tasks, the Dag processor checks each stub task
+against the handler registered by the artifact that a worker would run for it.
+A coordinator opts in by implementing two methods:
+
+.. code-block:: python
+
+    def _find_task_handler_artifact(self, *, bundle_path: pathlib.Path, dag_id: str) -> ResolvedBundle: ...
+
+
+    def _build_parse_task_handler_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: ...
+
+``_find_task_handler_artifact`` returns the artifact a stub task of *dag_id*
+runs and its supervisor schema version, found in the Dag bundle at
+*bundle_path* the same way ``_build_execute_task_command`` finds it. It scans
+*bundle_path* itself, because ``self._get_scan_roots()`` is not set there. It
+raises ``FileNotFoundError`` when the task would find no artifact it can run,
+or ``ValueError`` when the artifact it found uses a supervisor schema version
+this Task SDK does not know (``ExecutableCoordinator`` skips such bundles, so
+it raises ``FileNotFoundError`` instead).
+
+``_build_parse_task_handler_command`` builds the command that starts the
+runtime for the artifact at *path*, which ``_find_task_handler_artifact``
+returned. ``self._get_scan_roots()`` returns the root of the Dag bundle holding
+*path*, for a command that needs it, such as a classpath. The returned pair
+follows the rules of ``_build_execute_task_command``: no ``--comm`` or
+``--logs`` flags, and the schema version the runtime understands. A runtime
+whose schema version is older than ``TASK_HANDLER_PARSING_SCHEMA_VERSION`` is
+not started. The runtime answers as described in
+`Answering TaskHandlerParseRequest`_.
+
+Both defaults raise ``NotImplementedError``, so the stub tasks of a coordinator
+that does not implement them are not checked. ``ExecutableCoordinator``
+implements both for executable bundles.
+
 Supervisor Schema
 ~~~~~~~~~~~~~~~~~
 
@@ -195,8 +232,10 @@ as soon as possible. The supervisor verifies that the connecting peer belongs
 to the launched process tree, so the SDK MUST connect from the same process or
 one of its descendants.
 
-Once both connections are accepted, the supervisor sends a ``StartupDetails``
-message on the comm socket to initiate execution.
+Once both connections are accepted, the supervisor sends the first message on
+the comm socket. ``StartupDetails`` starts a task. A ``TaskHandlerParseRequest``
+comes from the Dag processor instead, and asks which task handlers the runtime
+registers (see `Answering TaskHandlerParseRequest`_).
 
 Wire protocol
 ~~~~~~~~~~~~~
@@ -273,6 +312,43 @@ calls themselves) can already produce log records before the ``--logs`` socket
 in stage 3 exists to carry them. See `Logging`_ below for how to handle that
 gap.
 
+Answering ``TaskHandlerParseRequest``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The SDK sends back one ``TaskHandlerParsingResult``, waits for the response to
+it, and exits. No task function runs: the answer lists the handlers
+registered through the SDK's ``TaskHandler`` interface, and never the tasks of
+a native Dag. ``task_handlers`` MUST declare every registered handler and
+depend only on the artifact, never on the request. ``fileloc`` repeats the
+request's ``file``.
+
+* ``task_handlers`` maps each Dag id the artifact registers handlers for to
+  their declarations, in registration order. A stub task the answer does not
+  declare makes the Dag fail to import.
+* Each ``TaskHandlerDeclaration`` states in ``binding`` how stub-task arguments
+  bind to its ``params``:
+
+  * ``positional``: by position. An argument count or value type the handler
+    cannot take makes the Dag fail to import.
+  * ``named``: by name in any order, ignoring case and underscores unless a
+    param sets ``exact_name``. An argument or param that matches nothing is
+    logged as a warning and the task still runs, so the runtime must accept
+    both. When no param matches and exactly one argument was passed, it may be
+    the whole value and is not warned about, unless ``params`` is empty, a
+    param sets ``exact_name``, or the argument cannot be an object. A value
+    type a param does not accept makes the Dag fail to import.
+
+* ``params`` lists the handler's parameters in order, ``[]`` when it has none,
+  or is ``null`` when the SDK cannot list them. Then only the handler's
+  presence is checked.
+* Each ``TaskHandlerParam`` has a ``name`` (``null`` when the SDK has no name
+  for a positional parameter) and a ``value_schema``: the JSON Schema of the
+  values it accepts, in the vocabulary ``@task.stub`` uses for Python
+  annotations, or ``null`` when the handler does not constrain it.
+* An empty ``task_handlers`` mapping is sent empty, never as ``null``.
+
+``go-sdk/pkg/execution`` is a reference implementation.
+
 Logging
 ~~~~~~~
 

go-sdk/README.md

@@ -389,6 +389,10 @@ Python supervisor / task runner
   protocol on the comm socket, with structured JSON-line logs on the logs socket.
 - The Python runtime is the worker. It proxies every `GetConnection` / `GetVariable` / `GetXCom` /
   `SetXCom` call through to the Execution API. The Go binary just runs the task function.
+- The Dag processor starts the binary a stub task would run in the same way, to check a Python Dag's stub
+  tasks. It sends a `TaskHandlerParseRequest`, and the binary replies with every task handler registered
+  with `airflow.TaskHandler`, keyed by Dag id in registration order, and how each handler's parameters
+  bind. Then it exits.
 
 The Go side of the protocol is implemented in `pkg/execution/`. On the Python side it is the
 `ExecutableCoordinator` in `task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py`.

go-sdk/airflow/bundle.go

@@ -130,18 +130,15 @@ func (b *BundleRef) Register(items ...Registerable) {
 }
 
 // taskHandlerMap holds the registered task handlers by dag_id and task_id.
-// It also keeps registration order. The --airflow-metadata manifest lists the tasks of each Dag
-// in that order.
+// It also keeps registration order. The --airflow-metadata manifest, and the reply to the Dag
+// processor's task handler parse, list the tasks of each Dag in that order.
 type taskHandlerMap struct {
 	mu       sync.RWMutex
 	handlers map[string]map[string]bundle.Task
 	order    []bundle.TaskHandlerInfo
 }
 
-var (
-	_ bundle.Bundle           = (*taskHandlerMap)(nil)
-	_ bundle.EnumerableBundle = (*taskHandlerMap)(nil)
-)
+var _ bundle.Registry = (*taskHandlerMap)(nil)
 
 func (m *taskHandlerMap) add(dagID, taskID string, task bundle.Task) {
 	m.mu.Lock()

go-sdk/airflow/dag.go

@@ -58,7 +58,8 @@ type DagRef struct {
 // panic once the Dag is registered.
 //
 // [BundleRef.Serve] does not yet serve the Dags that Dag returns. It leaves them out of the
-// --airflow-metadata manifest and cannot run their tasks.
+// --airflow-metadata manifest and cannot run their tasks. Their tasks are never task handlers, so
+// the reply to the Dag processor's task handler parse leaves them out too.
 func Dag(dagID string, spec ...DagSpec) *DagRef {
 	if len(spec) > 1 {
 		panic(fmt.Sprintf(

go-sdk/airflow/serve.go

@@ -57,8 +57,8 @@ const (
 // The command-line flags of the executable decide what Serve does.
 // With --airflow-metadata it prints the bundle's manifest and returns, which is how
 // airflow-go-pack reads the Dag and task ids of the registered task handlers.
-// With --comm and --logs, which the Airflow supervisor passes, it runs one task over the
-// coordinator protocol.
+// With --comm and --logs, which Airflow passes, it speaks the coordinator protocol: it either
+// runs one task, or tells the Dag processor every task handler registered with [TaskHandler].
 //
 // main must exit with a non-zero status when Serve returns an error, because the exit status
 // is how the supervisor learns that the task failed:

go-sdk/airflow/serve_test.go

@@ -284,3 +284,65 @@ func TestServeRunsTaskForSupervisor(t *testing.T) {
 	}
 	assert.True(t, ran)
 }
+
+// A fake Dag processor sends TaskHandlerParseRequest over the comm socket, as the probe does after
+// it starts the bundle with --comm and --logs. The tasks of a Dag from Dag are not task handlers,
+// so the reply leaves them out.
+func TestServeDeclaresTaskHandlersButNotDags(t *testing.T) {
+	commLn, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer commLn.Close()
+	logsLn, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer logsLn.Close()
+	deadline := time.Now().Add(10 * time.Second)
+	require.NoError(t, commLn.(*net.TCPListener).SetDeadline(deadline))
+	require.NoError(t, logsLn.(*net.TCPListener).SetDeadline(deadline))
+
+	native := Dag("native_etl")
+	native.Task(extract)
+	b := Bundle()
+	b.Register(TaskHandler("py_etl", "transform", noop), native)
+
+	done := make(chan error, 1)
+	go func() {
+		done <- b.serve(
+			[]string{"--comm", commLn.Addr().String(), "--logs", logsLn.Addr().String()},
+			io.Discard,
+		)
+	}()
+
+	commConn, err := commLn.Accept()
+	require.NoError(t, err)
+	defer commConn.Close()
+	logsConn, err := logsLn.Accept()
+	require.NoError(t, err)
+	defer logsConn.Close()
+	require.NoError(t, commConn.SetDeadline(deadline))
+
+	dagProcessor := execution.NewCoordinatorComm(commConn, commConn, discardLogger())
+	require.NoError(t, dagProcessor.SendRequest(0, map[string]any{
+		"type":        "TaskHandlerParseRequest",
+		"file":        "/bundles/etl",
+		"bundle_path": "/bundles",
+		"bundle_name": "go",
+	}))
+
+	frame, err := dagProcessor.ReadMessage()
+	require.NoError(t, err)
+	var body map[string]any
+	require.NoError(t, msgpack.Unmarshal(frame.Body, &body))
+	assert.Equal(t, "TaskHandlerParsingResult", body["type"])
+	handlers, ok := body["task_handlers"].(map[string]any)
+	require.True(t, ok)
+	assert.Len(t, handlers, 1)
+	assert.Contains(t, handlers, "py_etl")
+
+	require.NoError(t, dagProcessor.SendRequest(frame.ID, map[string]any{}))
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(2 * time.Second):
+		t.Fatal("Serve did not return after the result was acknowledged")
+	}
+}