From 57d20f71086632a9d6ffc34f81bfdd69529d22d5 Mon Sep 17 00:00:00 2001 From: henry3260 Date: Wed, 30 Sep 2026 23:18:39 +0800 Subject: [PATCH 1/2] Go SDK: bring the coordinator-protocol models back to the supervisor schema MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The models are generated from the supervisor wire-schema snapshot the Python Task SDK owns, and they had fallen two fields behind it: a snapshot can grow a property without the api_version moving, and the version constant's test is the only thing that was watching. Nothing else regenerated or verified the committed file, so the drift survived two releases in silence — msgpack drops a field the struct does not declare, so the supervisor sent multi_team and retry_reason and the Go SDK could not see them. The Dag specs had the same exposure and got a check for it. Making that check cover both generated surfaces is what keeps this from happening again, and is cheaper than watching one of them and hoping about the other. closes: #73954 --- .pre-commit-config.yaml | 12 +- go-sdk/README.md | 7 +- go-sdk/pkg/execution/genmodels/models.gen.go | 143 +++++++------- .../ci/prek/check_go_sdk_generated_drift.py | 182 ++++++++++++++++++ scripts/ci/prek/check_go_sdk_spec_drift.py | 140 -------------- ...y => test_check_go_sdk_generated_drift.py} | 72 +++++-- 6 files changed, 319 insertions(+), 237 deletions(-) create mode 100755 scripts/ci/prek/check_go_sdk_generated_drift.py delete mode 100755 scripts/ci/prek/check_go_sdk_spec_drift.py rename scripts/tests/ci/prek/{test_check_go_sdk_spec_drift.py => test_check_go_sdk_generated_drift.py} (58%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3b6fe5200e122..8657726fc527d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -348,21 +348,23 @@ repos: # prek run sync-java-sdk-dag-schema --hook-stage manual stages: ['manual'] files: ^java-sdk/sdk/schema/dag-schema\.json$ - - id: check-go-sdk-spec-drift - name: Check Go SDK Dag spec structs against the serialization schema - description: "Regenerate go-sdk/airflow/spec.gen.go and fail when the committed file differs" - entry: ./scripts/ci/prek/check_go_sdk_spec_drift.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 files: > (?x) ^airflow-core/src/airflow/serialization/schema\.json$| + ^task-sdk/src/airflow/sdk/execution_time/schema/schema\.json$| ^go-sdk/airflow/spec\.gen\.go$| ^go-sdk/airflow/spec\.go$| ^go-sdk/internal/genspec/.*\.go$| ^go-sdk/internal/genlicense/.*\.go$| - ^scripts/ci/prek/check_go_sdk_spec_drift\.py$ + ^go-sdk/pkg/execution/genmodels/.*\.go$| + ^scripts/ci/prek/check_go_sdk_generated_drift\.py$ pass_filenames: false require_serial: true - id: check-go-version-in-sync diff --git a/go-sdk/README.md b/go-sdk/README.md index c4b32712b2602..a6dedde08049b 100644 --- a/go-sdk/README.md +++ b/go-sdk/README.md @@ -391,7 +391,10 @@ SDK to a newer schema version: 2. Run `just generate-models`. `TestSupervisorSchemaVersionMatchesSnapshot` fails when the constant and the snapshot's `api_version` -drift, so a missed bump is caught by `go test` instead of needing a dedicated prek hook. +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 +when the committed files differ. ## Regenerating the Dag and task specs @@ -417,7 +420,7 @@ it into the authoring shape, each entry carrying the reason it exists: are hand-written in [`airflow/spec.go`](./airflow/spec.go); `TestTriggerRuleConstantsMatchPython` is their tripwire against Airflow's own enum. -The `check-go-sdk-spec-drift` prek hook regenerates the file and fails when the committed one +The `check-go-sdk-generated-drift` prek hook regenerates the file and fails when the committed one differs, so a schema change that never reached Go cannot merge. ## Architectural decisions diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 698f6c097f312..22741803d83e3 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -1624,6 +1624,9 @@ type TIRunContext struct { // MaxTries corresponds to the JSON schema field "max_tries". MaxTries int `msgpack:"max_tries"` + // MultiTeam corresponds to the JSON schema field "multi_team". + MultiTeam bool `msgpack:"multi_team,omitempty"` + // NextKwargs corresponds to the JSON schema field "next_kwargs". NextKwargs *NextKwargs `msgpack:"next_kwargs,omitempty"` @@ -1649,18 +1652,50 @@ type TIRunContext struct { type TaskArgBinding interface{} -type TriggerKwargs map[string]interface{} - -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` +type TaskBreadcrumbsResult struct { + // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". + Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` } -type Warnings []interface{} +type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} + +// Task callback status information. +// +// A Class with information about the success/failure TI callback to be executed. +// Currently, only failure +// callbacks when tasks are externally killed or experience heartbeat timeouts are +// run via DagFileProcessorProcess. +type TaskCallbackRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` + + // BundleVersion corresponds to the JSON schema field "bundle_version". + BundleVersion interface{} `msgpack:"bundle_version"` + + // ContextFromServer corresponds to the JSON schema field "context_from_server". + ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` + + // Filepath corresponds to the JSON schema field "filepath". + Filepath string `msgpack:"filepath"` + + // Msg corresponds to the JSON schema field "msg". + Msg interface{} `msgpack:"msg,omitempty"` + + // TaskCallbackType corresponds to the JSON schema field "task_callback_type". + TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` + + // TI corresponds to the JSON schema field "ti". + TI TaskInstance `msgpack:"ti"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` + + // VersionData corresponds to the JSON schema field "version_data". + VersionData *VersionData `msgpack:"version_data,omitempty"` +} type TaskIds []string @@ -1697,70 +1732,24 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } +type TaskInstanceState string + const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" const TaskInstanceStateDeferred TaskInstanceState = "deferred" const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateRemoved TaskInstanceState = "removed" const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateSuccess TaskInstanceState = "success" const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateQueued TaskInstanceState = "queued" const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" type TaskOutlets []AssetProfile -const TaskInstanceStateRemoved TaskInstanceState = "removed" - -type TaskInstanceState string - -type TaskBreadcrumbsResultBreadcrumbsElem map[string]interface{} - -type TaskBreadcrumbsResult struct { - // Breadcrumbs corresponds to the JSON schema field "breadcrumbs". - Breadcrumbs []TaskBreadcrumbsResultBreadcrumbsElem `msgpack:"breadcrumbs"` - - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` -} - -// Task callback status information. -// -// A Class with information about the success/failure TI callback to be executed. -// Currently, only failure -// callbacks when tasks are externally killed or experience heartbeat timeouts are -// run via DagFileProcessorProcess. -type TaskCallbackRequest struct { - // BundleName corresponds to the JSON schema field "bundle_name". - BundleName string `msgpack:"bundle_name"` - - // BundleVersion corresponds to the JSON schema field "bundle_version". - BundleVersion interface{} `msgpack:"bundle_version"` - - // ContextFromServer corresponds to the JSON schema field "context_from_server". - ContextFromServer *TIRunContext `msgpack:"context_from_server,omitempty"` - - // Filepath corresponds to the JSON schema field "filepath". - Filepath string `msgpack:"filepath"` - - // Msg corresponds to the JSON schema field "msg". - Msg interface{} `msgpack:"msg,omitempty"` - - // TaskCallbackType corresponds to the JSON schema field "task_callback_type". - TaskCallbackType interface{} `msgpack:"task_callback_type,omitempty"` - - // TI corresponds to the JSON schema field "ti". - TI TaskInstance `msgpack:"ti"` - - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` - - // VersionData corresponds to the JSON schema field "version_data". - VersionData *VersionData `msgpack:"version_data,omitempty"` -} - // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { // StartDate corresponds to the JSON schema field "start_date". @@ -1770,12 +1759,6 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateSkipped TaskStateState = "skipped" -const TaskStateStateRemoved TaskStateState = "removed" - // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1789,6 +1772,9 @@ type TaskState struct { // RenderedMapIndex corresponds to the JSON schema field "rendered_map_index". RenderedMapIndex interface{} `msgpack:"rendered_map_index,omitempty"` + // RetryReason corresponds to the JSON schema field "retry_reason". + RetryReason interface{} `msgpack:"retry_reason,omitempty"` + // State corresponds to the JSON schema field "state". State TaskStateState `msgpack:"state"` @@ -1796,6 +1782,12 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateRemoved TaskStateState = "removed" +const TaskStateStateSkipped TaskStateState = "skipped" + // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1845,7 +1837,7 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } -type VersionData map[string]interface{} +type TriggerKwargs map[string]interface{} // Update the response content part of an existing Human-in-the-loop response. type UpdateHITLDetail struct { @@ -1881,6 +1873,15 @@ type VariableKeysResult struct { Type string `msgpack:"type,omitempty"` } +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} + type VariableResult struct { // Key corresponds to the JSON schema field "key". Key string `msgpack:"key"` @@ -1892,6 +1893,10 @@ type VariableResult struct { Value interface{} `msgpack:"value"` } +type VersionData map[string]interface{} + +type Warnings []interface{} + // One positional stub-task argument pulled from an upstream task's XCom. type XComArgBinding struct { // Kind corresponds to the JSON schema field "kind". diff --git a/scripts/ci/prek/check_go_sdk_generated_drift.py b/scripts/ci/prek/check_go_sdk_generated_drift.py new file mode 100755 index 0000000000000..6133a0608c224 --- /dev/null +++ b/scripts/ci/prek/check_go_sdk_generated_drift.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +# 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. +""" +Keep the Go SDK's generated files in step with the schemas they generate from. + +Two of the Go SDK's surfaces are generated from schemas Python owns, and both are +committed, so nothing regenerates them when the schema on the Python side moves: + +* ``go-sdk/airflow/spec.gen.go`` — ``airflow.DagSpec`` and ``airflow.TaskSpec``, the + structs a Dag author fills in, from ``airflow-core``'s Dag serialization schema. +* ``go-sdk/pkg/execution/genmodels/`` — the coordinator-protocol messages, from the + supervisor wire-schema snapshot the Python Task SDK owns. + +Without this check a property added, renamed or retyped on the Python side leaves the +Go side silently behind: a Dag authored in Go keeps serializing the old shape, and +msgpack drops a message field the Go struct does not declare. That is not theoretical +— ``models.gen.go`` sat two fields behind its snapshot across two releases, which is +what #73954 is about. + +The check regenerates each target and asks Git whether it changed. A drifted file is +left regenerated in the working tree, so the fix is to commit it. + +Run from the repo root: + + uv run --project scripts python scripts/ci/prek/check_go_sdk_generated_drift.py + +Exits 0 if every committed file matches its schema, 1 otherwise. +""" + +from __future__ import annotations + +import os +import pathlib +import shutil +import subprocess +import sys +from typing import NamedTuple + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +GO_SDK_MODULE = pathlib.Path("go-sdk") + + +class Target(NamedTuple): + """One `go generate` target and the committed files it writes.""" + + package: str + committed: tuple[pathlib.Path, ...] + schema: str + # remedy is what to weigh before committing the regenerated file, or "" when there + # is nothing to decide and the file just has to be brought forward. + remedy: str + + +TARGETS = ( + Target( + package="./airflow/...", + committed=(GO_SDK_MODULE / "airflow" / "spec.gen.go",), + schema="airflow-core/src/airflow/serialization/schema.json", + remedy=( + "Review it — a property that should not reach an author belongs in the " + "exclusion list in go-sdk/internal/genspec/authoring.go — then commit it:" + ), + ), + Target( + package="./pkg/execution/genmodels/...", + committed=(GO_SDK_MODULE / "pkg" / "execution" / "genmodels",), + schema="task-sdk/src/airflow/sdk/execution_time/schema/schema.json", + remedy="Commit it:", + ), +) + + +def regenerate(module_dir: pathlib.Path, package: str, go_binary: str = "go") -> tuple[int, str]: + """Run one target's generators. Returns ``(returncode, combined_output)``.""" + completed = subprocess.run( + [go_binary, "generate", package], + cwd=module_dir, + capture_output=True, + text=True, + check=False, + ) + return completed.returncode, (completed.stdout + completed.stderr).strip() + + +def read_drift( + repo_root: pathlib.Path, paths: tuple[pathlib.Path, ...], git_binary: str = "git" +) -> tuple[int, str]: + """Ask Git what regeneration changed. Returns ``(returncode, diff)``.""" + completed = subprocess.run( + [git_binary, "diff", "--", *(str(path) for path in paths)], + cwd=repo_root, + capture_output=True, + text=True, + check=False, + ) + return completed.returncode, completed.stdout.strip() + + +def format_report( + target: Target, generate_returncode: int, generate_output: str, diff_returncode: int, diff: str +) -> tuple[int, str]: + """Turn one target's regeneration result and the diff that followed into ``(exit_code, report)``.""" + written = ", ".join(str(path) for path in target.committed) + if generate_returncode != 0: + return 1, "\n".join( + [ + f"ERROR: regenerating {written} failed.", + "", + "A generator fails on a schema construct it has no rule for, which is how a", + "change on the Python side that needs a new rule surfaces; the generators also", + "need the Go toolchain and the network to fetch go-jsonschema.", + "`go generate` reported:", + "", + generate_output or "(no output)", + ] + ) + # An unreadable diff is not an absent one: reporting success here would let the + # check pass on every drift. + if diff_returncode != 0: + return 1, f"ERROR: `git diff` failed, so whether {written} drifted is unknown." + if not diff: + return 0, f"OK: {written} matches {target.schema}." + return 1, "\n".join( + [ + f"ERROR: {written} is out of date.", + "", + f"It is generated from {target.schema}, which has moved since the file was", + "committed. The regenerated file is in your working tree.", + "", + target.remedy, + "", + f" git add {written}", + "", + "Regeneration changed:", + "", + diff, + ] + ) + + +def main() -> int: + module_dir = REPO_ROOT / GO_SDK_MODULE + for target in TARGETS: + for path in target.committed: + if not (REPO_ROOT / path).exists(): + print(f"ERROR: {path} not found — has the generated file moved?") + return 1 + if shutil.which("go") is None: + if os.environ.get("CI"): + print("ERROR: `go` is not on PATH but this is a CI run — the toolchain is required here.") + return 1 + print("SKIPPED: `go` is not on PATH, cannot verify that the generated files are current.") + return 0 + exit_code = 0 + for target in TARGETS: + generate_returncode, generate_output = regenerate(module_dir, target.package) + diff_returncode, diff = read_drift(REPO_ROOT, target.committed) + target_exit_code, report = format_report( + target, generate_returncode, generate_output, diff_returncode, diff + ) + print(report) + exit_code = exit_code or target_exit_code + return exit_code + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/ci/prek/check_go_sdk_spec_drift.py b/scripts/ci/prek/check_go_sdk_spec_drift.py deleted file mode 100755 index 9e5dfd0a3105a..0000000000000 --- a/scripts/ci/prek/check_go_sdk_spec_drift.py +++ /dev/null @@ -1,140 +0,0 @@ -#!/usr/bin/env python3 -# 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. -""" -Keep the Go SDK's Dag and task spec structs in step with the serialization schema. - -``go-sdk/airflow/spec.gen.go`` declares ``airflow.DagSpec`` and -``airflow.TaskSpec``, the two structs a Dag author fills in. Both are generated -from ``airflow-core/src/airflow/serialization/schema.json``, which Python owns: -``go-sdk/internal/genspec`` rewrites that schema into the authoring shape and -go-jsonschema writes the structs from it. - -The generated file is committed, so nothing regenerates it when the schema moves -on the Python side. Without this check a property added, renamed or retyped there -would leave the Go structs silently behind, and a Dag authored in Go would keep -serializing the old shape. - -The check regenerates the file and asks Git whether it changed. A drifted file is -left regenerated in the working tree, so the fix is to commit it. - -Run from the repo root: - - uv run --project scripts python scripts/ci/prek/check_go_sdk_spec_drift.py - -Exits 0 if the committed structs match the schema, 1 otherwise. -""" - -from __future__ import annotations - -import os -import pathlib -import shutil -import subprocess -import sys - -REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] -GO_SDK_MODULE = pathlib.Path("go-sdk") -GENERATED_SPECS = GO_SDK_MODULE / "airflow" / "spec.gen.go" - - -def regenerate_specs(module_dir: pathlib.Path, go_binary: str = "go") -> tuple[int, str]: - """Run the spec generators. Returns ``(returncode, combined_output)``.""" - completed = subprocess.run( - [go_binary, "generate", "./airflow/..."], - cwd=module_dir, - capture_output=True, - text=True, - check=False, - ) - return completed.returncode, (completed.stdout + completed.stderr).strip() - - -def read_drift(repo_root: pathlib.Path, git_binary: str = "git") -> tuple[int, str]: - """Ask Git what regeneration changed. Returns ``(returncode, diff)``.""" - completed = subprocess.run( - [git_binary, "diff", "--", str(GENERATED_SPECS)], - cwd=repo_root, - capture_output=True, - text=True, - check=False, - ) - return completed.returncode, completed.stdout.strip() - - -def format_report( - generate_returncode: int, generate_output: str, diff_returncode: int, diff: str -) -> tuple[int, str]: - """Turn a regeneration result and the diff that followed it into ``(exit_code, report)``.""" - if generate_returncode != 0: - return 1, "\n".join( - [ - f"ERROR: regenerating {GENERATED_SPECS} failed.", - "", - "genspec fails on a schema construct it has no rule for, which is how a change on", - "the Python side that needs a new rule surfaces; the generators also need the Go", - "toolchain and the network to fetch go-jsonschema. `go generate` reported:", - "", - generate_output or "(no output)", - ] - ) - # An unreadable diff is not an absent one: reporting success here would let the - # check pass on every drift. - if diff_returncode != 0: - return 1, f"ERROR: `git diff` failed, so whether {GENERATED_SPECS} drifted is unknown." - if not diff: - return 0, f"OK: {GENERATED_SPECS} matches the serialization schema." - return 1, "\n".join( - [ - f"ERROR: {GENERATED_SPECS} is out of date.", - "", - "It is generated from airflow-core/src/airflow/serialization/schema.json by", - "go-sdk/internal/genspec, and one of the two has moved since the file was committed.", - "The regenerated file is in your working tree.", - "", - "Review it — a property that should not reach an author belongs in the exclusion", - "list in go-sdk/internal/genspec/authoring.go — then commit it:", - "", - f" git add {GENERATED_SPECS}", - "", - "Regeneration changed:", - "", - diff, - ] - ) - - -def main() -> int: - module_dir = REPO_ROOT / GO_SDK_MODULE - if not (REPO_ROOT / GENERATED_SPECS).is_file(): - print(f"ERROR: {GENERATED_SPECS} not found — has the generated file moved?") - return 1 - if shutil.which("go") is None: - if os.environ.get("CI"): - print("ERROR: `go` is not on PATH but this is a CI run — the toolchain is required here.") - return 1 - print(f"SKIPPED: `go` is not on PATH, cannot verify that {GENERATED_SPECS} is current.") - return 0 - generate_returncode, generate_output = regenerate_specs(module_dir) - diff_returncode, diff = read_drift(REPO_ROOT) - exit_code, report = format_report(generate_returncode, generate_output, diff_returncode, diff) - print(report) - return exit_code - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/scripts/tests/ci/prek/test_check_go_sdk_spec_drift.py b/scripts/tests/ci/prek/test_check_go_sdk_generated_drift.py similarity index 58% rename from scripts/tests/ci/prek/test_check_go_sdk_spec_drift.py rename to scripts/tests/ci/prek/test_check_go_sdk_generated_drift.py index fffbd80c399a0..c1790faaddfe8 100644 --- a/scripts/tests/ci/prek/test_check_go_sdk_spec_drift.py +++ b/scripts/tests/ci/prek/test_check_go_sdk_generated_drift.py @@ -17,41 +17,71 @@ # under the License. from __future__ import annotations +import pathlib import subprocess from unittest import mock -import check_go_sdk_spec_drift as checker +import check_go_sdk_generated_drift as checker import pytest +SPECS, MODELS = checker.TARGETS + SPEC_DRIFT_DIFF = """\ diff --git a/go-sdk/airflow/spec.gen.go b/go-sdk/airflow/spec.gen.go --- a/go-sdk/airflow/spec.gen.go +++ b/go-sdk/airflow/spec.gen.go @@ -40,6 +40,9 @@ type DagSpec struct { + // Deadline corresponds to the JSON schema field "deadline". -+ Deadline string `json:"deadline,omitempty,omitzero"` ++ Deadline string +""" + +MODELS_DRIFT_DIFF = """\ +diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go +--- a/go-sdk/pkg/execution/genmodels/models.gen.go ++++ b/go-sdk/pkg/execution/genmodels/models.gen.go +@@ -1624,6 +1624,9 @@ type TIRunContext struct { ++ MultiTeam bool `msgpack:"multi_team,omitempty"` """ -def test_current_specs_pass(): - exit_code, report = checker.format_report(0, "", 0, "") +def test_both_targets_are_checked(): + assert [target.package for target in checker.TARGETS] == [ + "./airflow/...", + "./pkg/execution/genmodels/...", + ] + + +def test_current_files_pass(): + exit_code, report = checker.format_report(SPECS, 0, "", 0, "") assert exit_code == 0 - assert "matches the serialization schema" in report + assert "matches airflow-core/src/airflow/serialization/schema.json" in report def test_drifted_specs_fail_with_the_diff_and_where_to_decide_about_a_property(): - exit_code, report = checker.format_report(0, "", 0, SPEC_DRIFT_DIFF) + exit_code, report = checker.format_report(SPECS, 0, "", 0, SPEC_DRIFT_DIFF) assert exit_code == 1 assert "is out of date" in report assert "go-sdk/internal/genspec/authoring.go" in report assert "git add go-sdk/airflow/spec.gen.go" in report - assert 'Deadline string `json:"deadline,omitempty,omitzero"`' in report + assert "Deadline string" in report + + +def test_drifted_models_name_the_supervisor_snapshot_and_have_nothing_to_decide(): + exit_code, report = checker.format_report(MODELS, 0, "", 0, MODELS_DRIFT_DIFF) + + assert exit_code == 1 + assert "task-sdk/src/airflow/sdk/execution_time/schema/schema.json" in report + assert "git add go-sdk/pkg/execution/genmodels" in report + # Nothing is excluded from the models, so there is no list to weigh a field against. + assert "authoring.go" not in report + assert "MultiTeam" in report def test_failed_generation_reports_the_generator_output_instead_of_a_diff(): exit_code, report = checker.format_report( + SPECS, 1, "genspec: shaping schema.json for authoring: definitions/dag/properties/fileloc is excluded", 0, @@ -64,47 +94,47 @@ def test_failed_generation_reports_the_generator_output_instead_of_a_diff(): def test_failed_generation_without_output_still_reports(): - exit_code, report = checker.format_report(1, "", 0, "") + exit_code, report = checker.format_report(SPECS, 1, "", 0, "") assert exit_code == 1 assert "(no output)" in report def test_unreadable_diff_fails_instead_of_passing_as_no_drift(): - exit_code, report = checker.format_report(0, "", 128, "") + exit_code, report = checker.format_report(SPECS, 0, "", 128, "") assert exit_code == 1 assert "is unknown" in report -@mock.patch("check_go_sdk_spec_drift.subprocess.run", autospec=True) -def test_regeneration_runs_the_generators_in_the_go_sdk_module(mock_run, tmp_path): +@mock.patch("check_go_sdk_generated_drift.subprocess.run", autospec=True) +def test_regeneration_runs_one_targets_generators_in_the_go_sdk_module(mock_run, tmp_path): mock_run.return_value = subprocess.CompletedProcess(args=[], returncode=0, stdout="", stderr="") - checker.regenerate_specs(tmp_path) + checker.regenerate(tmp_path, MODELS.package) - assert mock_run.call_args.args[0] == ["go", "generate", "./airflow/..."] + assert mock_run.call_args.args[0] == ["go", "generate", "./pkg/execution/genmodels/..."] assert mock_run.call_args.kwargs["cwd"] == tmp_path -@mock.patch("check_go_sdk_spec_drift.subprocess.run", autospec=True) +@mock.patch("check_go_sdk_generated_drift.subprocess.run", autospec=True) def test_regeneration_combines_stdout_and_stderr(mock_run, tmp_path): mock_run.return_value = subprocess.CompletedProcess( args=[], returncode=1, stdout="genspec: shaping failed\n", stderr="exit status 1\n" ) - returncode, output = checker.regenerate_specs(tmp_path) + returncode, output = checker.regenerate(tmp_path, SPECS.package) assert returncode == 1 assert "genspec: shaping failed" in output assert "exit status 1" in output -@mock.patch("check_go_sdk_spec_drift.subprocess.run", autospec=True) -def test_read_drift_asks_git_only_about_the_generated_file(mock_run, tmp_path): +@mock.patch("check_go_sdk_generated_drift.subprocess.run", autospec=True) +def test_read_drift_asks_git_only_about_one_targets_files(mock_run, tmp_path): mock_run.return_value = subprocess.CompletedProcess(args=[], returncode=0, stdout="", stderr="") - checker.read_drift(tmp_path) + checker.read_drift(tmp_path, (pathlib.Path("go-sdk/airflow/spec.gen.go"),)) assert mock_run.call_args.args[0] == ["git", "diff", "--", "go-sdk/airflow/spec.gen.go"] assert mock_run.call_args.kwargs["cwd"] == tmp_path @@ -117,7 +147,7 @@ def test_read_drift_asks_git_only_about_the_generated_file(mock_run, tmp_path): pytest.param({}, 0, "SKIPPED", id="local-skips"), ], ) -@mock.patch("check_go_sdk_spec_drift.shutil.which", autospec=True, return_value=None) +@mock.patch("check_go_sdk_generated_drift.shutil.which", autospec=True, return_value=None) def test_missing_go_toolchain(mock_which, ci_env, expected_exit, expected_text, monkeypatch, capsys): monkeypatch.delenv("CI", raising=False) for key, value in ci_env.items(): @@ -127,7 +157,7 @@ def test_missing_go_toolchain(mock_which, ci_env, expected_exit, expected_text, assert expected_text in capsys.readouterr().out -@mock.patch("check_go_sdk_spec_drift.pathlib.Path.is_file", autospec=True, return_value=False) -def test_missing_generated_file_fails(mock_is_file, capsys): +@mock.patch("check_go_sdk_generated_drift.pathlib.Path.exists", autospec=True, return_value=False) +def test_missing_generated_file_fails(mock_exists, capsys): assert checker.main() == 1 assert "not found" in capsys.readouterr().out From 52440b1ef52acfb586e642f6c4b3c692f6d12f2c Mon Sep 17 00:00:00 2001 From: henry3260 Date: Wed, 30 Sep 2026 23:57:56 +0800 Subject: [PATCH 2/2] Go SDK: name the generated model files the drift check watches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The models target pointed `git diff` at the whole genmodels package, but that package also holds hand-written code: gen.go carries the //go:generate directives, gen/ the custom generator, plus two test files. The hook's own files: regex fires on edits to exactly those, so editing a generator made the check report the author's own uncommitted gen.go edit as generator drift, and told them to `git add` the package directory, staging unrelated work with it. A directory also defeats the missing-file guard, which the specs target gets for free from naming a file: a directory outlives every generated file in it, so a moved or deleted models.gen.go passed silently — the drift the guard exists to catch. Naming the files also lets a report say which one drifted rather than which package, which is the one thing it has to get across. Three of them do not fit the singular wording the reports had, and a comma-joined list is not a runnable `git add`, so both follow from the same change. --- .../ci/prek/check_go_sdk_generated_drift.py | 21 ++++++++++++------- .../prek/test_check_go_sdk_generated_drift.py | 17 +++++++++++---- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/scripts/ci/prek/check_go_sdk_generated_drift.py b/scripts/ci/prek/check_go_sdk_generated_drift.py index 6133a0608c224..f66dcae71b9ee 100755 --- a/scripts/ci/prek/check_go_sdk_generated_drift.py +++ b/scripts/ci/prek/check_go_sdk_generated_drift.py @@ -23,8 +23,8 @@ * ``go-sdk/airflow/spec.gen.go`` — ``airflow.DagSpec`` and ``airflow.TaskSpec``, the structs a Dag author fills in, from ``airflow-core``'s Dag serialization schema. -* ``go-sdk/pkg/execution/genmodels/`` — the coordinator-protocol messages, from the - supervisor wire-schema snapshot the Python Task SDK owns. +* ``go-sdk/pkg/execution/genmodels/*.gen.go`` — the coordinator-protocol messages, + from the supervisor wire-schema snapshot the Python Task SDK owns. Without this check a property added, renamed or retyped on the Python side leaves the Go side silently behind: a Dag authored in Go keeps serializing the old shape, and @@ -78,7 +78,11 @@ class Target(NamedTuple): ), Target( package="./pkg/execution/genmodels/...", - committed=(GO_SDK_MODULE / "pkg" / "execution" / "genmodels",), + committed=( + GO_SDK_MODULE / "pkg" / "execution" / "genmodels" / "models.gen.go", + GO_SDK_MODULE / "pkg" / "execution" / "genmodels" / "discriminators.gen.go", + GO_SDK_MODULE / "pkg" / "execution" / "genmodels" / "defaults.gen.go", + ), schema="task-sdk/src/airflow/sdk/execution_time/schema/schema.json", remedy="Commit it:", ), @@ -116,6 +120,7 @@ def format_report( ) -> tuple[int, str]: """Turn one target's regeneration result and the diff that followed into ``(exit_code, report)``.""" written = ", ".join(str(path) for path in target.committed) + stageable = " ".join(str(path) for path in target.committed) if generate_returncode != 0: return 1, "\n".join( [ @@ -134,17 +139,17 @@ def format_report( if diff_returncode != 0: return 1, f"ERROR: `git diff` failed, so whether {written} drifted is unknown." if not diff: - return 0, f"OK: {written} matches {target.schema}." + return 0, f"OK: {written} — up to date with {target.schema}." return 1, "\n".join( [ - f"ERROR: {written} is out of date.", + f"ERROR: out of date: {written}.", "", - f"It is generated from {target.schema}, which has moved since the file was", - "committed. The regenerated file is in your working tree.", + f"The committed output no longer matches what {target.schema} and the", + "generators produce. The regenerated output is in your working tree.", "", target.remedy, "", - f" git add {written}", + f" git add {stageable}", "", "Regeneration changed:", "", diff --git a/scripts/tests/ci/prek/test_check_go_sdk_generated_drift.py b/scripts/tests/ci/prek/test_check_go_sdk_generated_drift.py index c1790faaddfe8..de4efb557885a 100644 --- a/scripts/tests/ci/prek/test_check_go_sdk_generated_drift.py +++ b/scripts/tests/ci/prek/test_check_go_sdk_generated_drift.py @@ -44,25 +44,30 @@ """ -def test_both_targets_are_checked(): +def test_both_targets_are_checked_and_name_generated_files_not_a_package_directory(): assert [target.package for target in checker.TARGETS] == [ "./airflow/...", "./pkg/execution/genmodels/...", ] + # A package directory would widen `git diff` onto the hand-written gen.go beside the + # generated files, and hide a deleted one from the missing-file guard in main(). + for target in checker.TARGETS: + assert target.committed + assert all(path.name.endswith(".gen.go") for path in target.committed), target.committed def test_current_files_pass(): exit_code, report = checker.format_report(SPECS, 0, "", 0, "") assert exit_code == 0 - assert "matches airflow-core/src/airflow/serialization/schema.json" in report + assert "up to date with airflow-core/src/airflow/serialization/schema.json" in report def test_drifted_specs_fail_with_the_diff_and_where_to_decide_about_a_property(): exit_code, report = checker.format_report(SPECS, 0, "", 0, SPEC_DRIFT_DIFF) assert exit_code == 1 - assert "is out of date" in report + assert "out of date" in report assert "go-sdk/internal/genspec/authoring.go" in report assert "git add go-sdk/airflow/spec.gen.go" in report assert "Deadline string" in report @@ -73,7 +78,11 @@ def test_drifted_models_name_the_supervisor_snapshot_and_have_nothing_to_decide( assert exit_code == 1 assert "task-sdk/src/airflow/sdk/execution_time/schema/schema.json" in report - assert "git add go-sdk/pkg/execution/genmodels" in report + assert ( + "git add go-sdk/pkg/execution/genmodels/models.gen.go " + "go-sdk/pkg/execution/genmodels/discriminators.gen.go " + "go-sdk/pkg/execution/genmodels/defaults.gen.go" in report + ) # Nothing is excluded from the models, so there is no list to weigh a field against. assert "authoring.go" not in report assert "MultiTeam" in report