Go SDK: vendor the schemas it generates from (#74047)

Author: henry3260

.pre-commit-config.yaml

@@ -348,17 +348,38 @@ repos:
         #   prek run sync-java-sdk-dag-schema --hook-stage manual
         stages: ['manual']
         files: ^java-sdk/sdk/schema/dag-schema\.json$
+      - id: sync-go-sdk-schemas
+        name: Sync Go SDK vendored schemas with the distributions that own them
+        description: "Copy airflow-core's Dag schema and task-sdk's supervisor schema when go-sdk's vendored copies drift"
+        entry: ./scripts/ci/prek/sync_go_sdk_schemas.py
+        language: python
+        pass_filenames: false
+        # Re-vendoring is a go-sdk maintainer's deliberate step after a schema change on
+        # the Python side, not something every commit should do, so this is manual only,
+        # the way sync-java-sdk-dag-schema above is:
+        #   prek run sync-go-sdk-schemas --hook-stage manual
+        # Running it on every commit would make a change to a Python schema fail the PR
+        # of whoever made it, which is the coupling vendoring is here to remove.
+        stages: ['manual']
+        files: >
+          (?x)
+          ^airflow-core/src/airflow/serialization/schema\.json$|
+          ^task-sdk/src/airflow/sdk/execution_time/schema/schema\.json$|
+          ^go-sdk/schema/.*\.json$|
+          ^scripts/ci/prek/sync_go_sdk_schemas\.py$
       - id: check-go-sdk-generated-drift
         name: Check Go SDK generated files against the schemas they generate from
         description: "Regenerate the Dag specs and the coordinator-protocol models, and fail when a committed file differs"
         entry: ./scripts/ci/prek/check_go_sdk_generated_drift.py
         # golang so prek provisions the toolchain the generators need, the way the
         # other checks that shell out to `go` in go-sdk get theirs.
         language: golang
+        # The vendored schemas, not the Python originals: sync-go-sdk-schemas above is
+        # what ties those to the copies, as a manual step. This one runs on every commit,
+        # because a refreshed copy that nothing regenerated from is the drift that matters.
         files: >
           (?x)
-          ^airflow-core/src/airflow/serialization/schema\.json$|
-          ^task-sdk/src/airflow/sdk/execution_time/schema/schema\.json$|
+          ^go-sdk/schema/.*\.json$|
           ^go-sdk/airflow/spec\.gen\.go$|
           ^go-sdk/airflow/spec\.go$|
           ^go-sdk/internal/genspec/.*\.go$|

go-sdk/README.md

@@ -380,17 +380,36 @@ Python supervisor / task runner
 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`.
 
+## The vendored schemas
+
+`schema/dag-schema.json` and `schema/supervisor-schema.json` are byte-for-byte copies of schemas
+airflow-core and task-sdk own. The generators read the copies, not the originals, so a generated file
+is explainable from a file inside this module and a standalone checkout — a published Go module, an
+ASF source release — can regenerate and test without the monorepo around it. `ts-sdk` and `java-sdk`
+vendor theirs the same way.
+
+Two prek hooks split the work that keeps the copies honest. `sync-go-sdk-schemas` copies an original
+over its copy when the two differ and fails, since copying is mechanical; it is a **manual** hook, so
+re-vendoring is a deliberate step here rather than something a change to a Python schema forces on
+whoever made it — which also means nothing announces on its own that a copy has gone stale.
+`check-go-sdk-generated-drift` runs on every commit and regenerates from the copies, failing when the
+committed Go differs, since what to do about a new schema construct — a generator rule, an authoring
+exclusion — is a decision.
+
 ## Regenerating the coordinator-protocol models
 
-The types in `pkg/execution/genmodels/` are generated from the in-tree supervisor schema snapshot
-(`task-sdk/src/airflow/sdk/execution_time/schema/schema.json`); do not edit them by hand. To move the
-SDK to a newer schema version:
+The types in `pkg/execution/genmodels/` are generated from `schema/supervisor-schema.json`, this
+module's vendored copy of the supervisor schema snapshot the Python Task SDK owns
+(`task-sdk/src/airflow/sdk/execution_time/schema/schema.json`); do not edit either by hand. To move
+the SDK to a newer schema version:
 
-1. Set `SupervisorSchemaVersion` in [`pkg/execution/messages.go`](./pkg/execution/messages.go) to the
-   snapshot's `api_version` date.
-2. Run `just generate-models`.
+1. Refresh the copy: `prek run sync-go-sdk-schemas --hook-stage manual`, which overwrites
+   `schema/supervisor-schema.json` from task-sdk's and fails so the change lands in review.
+2. Set `SupervisorSchemaVersion` in [`pkg/execution/messages.go`](./pkg/execution/messages.go) to the
+   copy's `api_version` date.
+3. Run `just generate-models`, and commit the copy with what it generated.
 
-`TestSupervisorSchemaVersionMatchesSnapshot` fails when the constant and the snapshot's `api_version`
+`TestSupervisorSchemaVersionMatchesSnapshot` fails when the constant and the copy's `api_version`
 drift, so a missed bump is caught by `go test`. A snapshot can also grow a field without the
 `api_version` moving, which leaves the models behind with nothing failing — msgpack drops a field the
 struct does not declare — so the `check-go-sdk-generated-drift` prek hook regenerates them and fails
@@ -404,9 +423,10 @@ the models adds, renames or removes a `DagRunState` constant. The test keeps fai
 ## Regenerating the Dag and task specs
 
 `airflow.DagSpec` and `airflow.TaskSpec` in [`airflow/spec.gen.go`](./airflow/spec.gen.go) are
-generated from airflow-core's Dag serialization schema
-(`airflow-core/src/airflow/serialization/schema.json`), which Python owns; do not edit them by hand.
-Run `just generate-specs` after changing the schema or the generator.
+generated from `schema/dag-schema.json`, this module's vendored copy of airflow-core's Dag
+serialization schema (`airflow-core/src/airflow/serialization/schema.json`), which Python owns; do
+not edit either by hand. Refresh the copy with `prek run sync-go-sdk-schemas --hook-stage manual`,
+then run `just generate-specs` after changing the schema or the generator.
 
 The schema is the serialized shape rather than the authoring one, so
 [`internal/genspec/authoring.go`](./internal/genspec/authoring.go) holds the three tables that turn

go-sdk/airflow/spec.go

@@ -35,7 +35,7 @@ import "reflect"
 // the zero value. Serializing a Dag converts the fields; a tag would invite
 // json.Marshal(spec) to skip that step and produce a shape core misreads.
 
-//go:generate go run ../internal/genspec -schema ../../airflow-core/src/airflow/serialization/schema.json -out ../../.build/go-sdk/spec.schema.json
+//go:generate go run ../internal/genspec -schema ../schema/dag-schema.json -out ../../.build/go-sdk/spec.schema.json
 //go:generate go run github.com/atombender/go-jsonschema@v0.23.1 --only-models --struct-name-from-title --tags "" --capitalization ID --capitalization JSON --capitalization MD --capitalization XCom -p airflow -o spec.gen.go ../../.build/go-sdk/spec.schema.json
 //go:generate go run ../internal/genspec -license spec.gen.go
 

go-sdk/internal/genspec/main.go

@@ -15,10 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
-// Command genspec rewrites Airflow core's Dag serialization schema
-// (airflow-core/src/airflow/serialization/schema.json) into the schema the airflow
-// package's DagSpec and TaskSpec generate from, so that the two structs are not
-// hand-maintained.
+// Command genspec rewrites Airflow core's Dag serialization schema (go-sdk's vendored
+// copy of it, schema/dag-schema.json) into the schema the airflow package's DagSpec
+// and TaskSpec generate from, so that the two structs are not hand-maintained.
 //
 // The schema is owned by Python and stays untouched; the rewritten copy is a build
 // artifact. genspec rewrites it in two passes: shapeForAuthoring turns the
@@ -45,7 +44,7 @@ func main() {
 	schemaPath := flag.String(
 		"schema",
 		"",
-		"path to airflow-core's Dag serialization schema.json",
+		"path to the vendored copy of airflow-core's Dag serialization schema",
 	)
 	outPath := flag.String("out", "", "path to write the rewritten schema to")
 	licensePath := flag.String(

go-sdk/internal/genspec/normalize_test.go

@@ -26,8 +26,9 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// coreSchemaPath is the schema genspec normalizes, reached from this package.
-const coreSchemaPath = "../../../airflow-core/src/airflow/serialization/schema.json"
+// coreSchemaPath is the schema genspec normalizes: go-sdk's vendored copy of
+// airflow-core's, so this test runs in a standalone checkout as well.
+const coreSchemaPath = "../../schema/dag-schema.json"
 
 func schemaFrom(t *testing.T, body string) map[string]any {
 	t.Helper()

go-sdk/pkg/execution/genmodels/gen.go

@@ -19,13 +19,14 @@
 // messages (msgpack-over-socket) exchanged between the Airflow supervisor and a
 // Go SDK bundle in coordinator mode.
 //
-// These files are generated from the supervisor wire-schema snapshot owned by the
-// Python Task SDK (task-sdk/src/airflow/sdk/execution_time/schema/schema.json):
+// These files are generated from ../../../schema/supervisor-schema.json, go-sdk's
+// vendored copy of the supervisor wire-schema snapshot the Python Task SDK owns:
 // models.gen.go holds the struct types, discriminators.gen.go the Type<Name>
 // constants and EnsureType, defaults.gen.go the DecodeMsgpack methods that seed
 // non-zero schema defaults. Don't edit them by hand; change the Pydantic models,
-// let the generate-supervisor-schemas-snapshot prek hook refresh the snapshot,
-// then re-run `just generate-models` (go generate).
+// let the generate-supervisor-schemas-snapshot prek hook refresh the snapshot and
+// sync-go-sdk-schemas refresh the vendored copy, then re-run `just generate-models`
+// (go generate).
 //
 // go generate runs go-jsonschema, then the local gen tool, which strips
 // go-jsonschema's dead anyOf-branch typedefs, widens concrete int/float/bool
@@ -35,5 +36,5 @@
 // discriminators.gen.go and defaults.gen.go.
 package genmodels
 
-//go:generate go run github.com/atombender/go-jsonschema@v0.23.1 --only-models --struct-name-from-title --tags msgpack --capitalization ID --capitalization URI --capitalization TI -p genmodels -o models.gen.go ../../../../task-sdk/src/airflow/sdk/execution_time/schema/schema.json
-//go:generate go run ./gen -schema ../../../../task-sdk/src/airflow/sdk/execution_time/schema/schema.json -models models.gen.go -out discriminators.gen.go -defaults defaults.gen.go
+//go:generate go run github.com/atombender/go-jsonschema@v0.23.1 --only-models --struct-name-from-title --tags msgpack --capitalization ID --capitalization URI --capitalization TI -p genmodels -o models.gen.go ../../../schema/supervisor-schema.json
+//go:generate go run ./gen -schema ../../../schema/supervisor-schema.json -models models.gen.go -out discriminators.gen.go -defaults defaults.gen.go

go-sdk/pkg/execution/messages_test.go

@@ -30,25 +30,18 @@ import (
 	"github.com/apache/airflow/go-sdk/pkg/execution/genmodels"
 )
 
-// sourceSchemaRelPath locates the canonical supervisor schema snapshot owned by
-// the Python Task SDK, relative to this package directory within the monorepo.
-// It is the same file genmodels is generated against (referenced by relative
-// path rather than vendored into go-sdk).
-const sourceSchemaRelPath = "../../../task-sdk/src/airflow/sdk/execution_time/schema/schema.json"
+// vendoredSchemaRelPath locates go-sdk's vendored copy of the supervisor schema
+// snapshot the Python Task SDK owns, relative to this package directory. It is
+// the file genmodels is generated against, and it ships with go-sdk, so this
+// test needs no monorepo around it. The sync-go-sdk-schemas prek hook is what
+// keeps the copy equal to task-sdk's.
+const vendoredSchemaRelPath = "../../schema/supervisor-schema.json"
 
 // TestSupervisorSchemaVersionMatchesSnapshot pins SupervisorSchemaVersion to the
 // api_version of the schema the models are generated from, so a schema bump that
-// forgets to update the constant (or vice versa) fails loudly. It only runs when
-// the task-sdk source is reachable from the checkout, so a standalone go-sdk
-// build (without the rest of the monorepo) does not fail.
+// forgets to update the constant (or vice versa) fails loudly.
 func TestSupervisorSchemaVersionMatchesSnapshot(t *testing.T) {
-	raw, err := os.ReadFile(sourceSchemaRelPath)
-	if os.IsNotExist(err) {
-		t.Skipf(
-			"task-sdk schema source not reachable at %s; skipping version check",
-			sourceSchemaRelPath,
-		)
-	}
+	raw, err := os.ReadFile(vendoredSchemaRelPath)
 	require.NoError(t, err)
 	var snapshot struct {
 		APIVersion string `json:"api_version"`