AIR CLI Integration: --override flag functionality

acceptance/experimental/air/run/output.txt

@@ -14,9 +14,22 @@ Dry run: configuration for "smoke-test" is valid; not submitting.
   }
 }
 
-=== override not yet supported
->>> [CLI] experimental air run -f valid.yaml --dry-run --override a=b
-Error: --override is not yet supported
+=== override applies and logs the change
+>>> [CLI] experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=2 --override timeout_minutes=45
+Override: changing compute.num_accelerators from 1 to 2
+Override: setting timeout_minutes to 45
+Dry run: configuration for "smoke-test" is valid; not submitting.
+
+=== override of an unknown field is rejected
+>>> [CLI] experimental air run -f valid.yaml --dry-run --override bogus=1
+Error: invalid --override "bogus": "bogus" is not a known field; available fields are: code_source, command, compute, env_variables, environment, experiment_name, idempotency_token, max_retries, mlflow_experiment_directory, mlflow_run_name, parameters, permissions, secrets, timeout_minutes, usage_policy_id, usage_policy_name
+
+Exit code: 1
+
+=== override still runs schema validation
+>>> [CLI] experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0
+Override: changing compute.num_accelerators from 1 to 0
+Error: compute.num_accelerators must be positive, got 0
 
 Exit code: 1
 

acceptance/experimental/air/run/script

@@ -4,8 +4,14 @@ trace $CLI experimental air run -f valid.yaml --dry-run
 title "dry-run (json)"
 trace $CLI experimental air run -f valid.yaml --dry-run -o json
 
-title "override not yet supported"
-errcode trace $CLI experimental air run -f valid.yaml --dry-run --override a=b
+title "override applies and logs the change"
+trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=2 --override timeout_minutes=45
+
+title "override of an unknown field is rejected"
+errcode trace $CLI experimental air run -f valid.yaml --dry-run --override bogus=1
+
+title "override still runs schema validation"
+errcode trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0
 
 title "watch not yet supported"
 errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch

experimental/air/cmd/run.go

@@ -57,16 +57,12 @@ The workload is described by a YAML config file (see --file).`,
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
 
-		// These flags' pipelines are not ported yet; reject rather than silently
-		// ignore them.
-		if len(overrides) > 0 {
-			return errors.New("--override is not yet supported")
-		}
+		// --watch's pipeline is not ported yet; reject rather than silently ignore it.
 		if watch {
 			return errors.New("--watch is not yet supported")
 		}
 
-		cfg, err := loadRunConfig(file)
+		cfg, err := loadRunConfigWithOverrides(ctx, file, overrides)
 		if err != nil {
 			return err
 		}

experimental/air/cmd/runconfig_load.go

@@ -1,6 +1,8 @@
 package aircmd
 
 import (
+	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -10,18 +12,25 @@ import (
 )
 
 // decodeRunConfig reads and decodes the run YAML into the schema. Unknown keys
-// are rejected (KnownFields), mirroring the Python schema's extra="forbid".
+// are rejected (KnownFields).
 //
-// The `_bases_` composition feature and CLI `--override` handling are not yet
-// ported; a config using `_bases_` is currently rejected as an unknown field.
+// The `_bases_` composition feature is not yet ported; a config using `_bases_`
+// is currently rejected as an unknown field. CLI `--override` handling lives in
+// runconfig_override.go and is applied to the parsed map before this decode.
 func decodeRunConfig(path string) (*runConfig, error) {
 	f, err := os.Open(path)
 	if err != nil {
 		return nil, err
 	}
 	defer f.Close()
 
-	dec := yaml.NewDecoder(f)
+	return decodeRunConfigReader(f, path)
+}
+
+// decodeRunConfigReader decodes and unknown-key-checks a run YAML from r. path is
+// used only for error messages.
+func decodeRunConfigReader(r io.Reader, path string) (*runConfig, error) {
+	dec := yaml.NewDecoder(r)
 	dec.KnownFields(true)
 
 	var cfg runConfig
@@ -50,3 +59,53 @@ func loadRunConfig(path string) (*runConfig, error) {
 	}
 	return cfg, nil
 }
+
+// loadRunConfigWithOverrides decodes a run YAML config, applies any
+// --override KEY=VALUE entries to the parsed map, then re-decodes (with unknown
+// keys rejected) and structurally validates the result. Applying overrides to
+// the map — rather than the typed config — lets the single decode+validate
+// pipeline enforce path existence, type coercion, and the semantic rules at
+// once. ctx is used only to log applied overrides.
+func loadRunConfigWithOverrides(ctx context.Context, path string, overrides []string) (*runConfig, error) {
+	if len(overrides) == 0 {
+		return loadRunConfig(path)
+	}
+
+	entries, err := parseOverrides(overrides)
+	if err != nil {
+		return nil, err
+	}
+	if err := validateOverridePaths(entries); err != nil {
+		return nil, err
+	}
+
+	raw, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+	var m map[string]any
+	if err := yaml.Unmarshal(raw, &m); err != nil {
+		return nil, fmt.Errorf("invalid config %s: %w", path, err)
+	}
+	if m == nil {
+		// An empty file decodes to a nil map; start from an empty one so overrides
+		// can populate it (the re-decode still enforces required fields).
+		m = map[string]any{}
+	}
+	if err := applyOverrides(ctx, m, entries); err != nil {
+		return nil, err
+	}
+
+	merged, err := yaml.Marshal(m)
+	if err != nil {
+		return nil, err
+	}
+	cfg, err := decodeRunConfigReader(bytes.NewReader(merged), path)
+	if err != nil {
+		return nil, err
+	}
+	if err := validateRunConfig(cfg); err != nil {
+		return nil, err
+	}
+	return cfg, nil
+}

experimental/air/cmd/runconfig_override.go

@@ -0,0 +1,162 @@
+package aircmd
+
+import (
+	"context"
+	"fmt"
+	"maps"
+	"reflect"
+	"slices"
+	"strings"
+
+	"github.com/databricks/cli/libs/cmdio"
+	"go.yaml.in/yaml/v3"
+)
+
+// This file implements the `--override KEY=VALUE` flag. Overrides are applied to
+// the parsed YAML map (not the typed runConfig) before re-decode, so one pipeline
+// covers path existence, type coercion, and the semantic validate() rules.
+
+// freeFormFields hold free-form maps, so path validation stops at them: any
+// sub-path is valid.
+var freeFormFields = map[string]bool{
+	"parameters":    true,
+	"env_variables": true,
+	"secrets":       true,
+}
+
+// parseOverrides parses --override KEY=VALUE arguments, preserving order.
+func parseOverrides(overrides []string) ([]overrideEntry, error) {
+	entries := make([]overrideEntry, 0, len(overrides))
+	for _, item := range overrides {
+		key, value, found := strings.Cut(item, "=")
+		if !found {
+			// --override is repeatable, so a config path meant for -f can be
+			// swallowed here; point at the real fix.
+			hint := ""
+			if strings.HasSuffix(item, ".yaml") || strings.HasSuffix(item, ".yml") {
+				hint = fmt.Sprintf("; %q looks like a config file — pass it with -f/--file", item)
+			}
+			return nil, fmt.Errorf("invalid --override %q: expected KEY=VALUE (e.g. compute.num_accelerators=32)%s", item, hint)
+		}
+		key = strings.TrimSpace(key)
+		if key == "" {
+			return nil, fmt.Errorf("invalid --override %q: empty key", item)
+		}
+		entries = append(entries, overrideEntry{path: key, raw: value})
+	}
+	return entries, nil
+}
+
+// overrideEntry is one parsed --override: its dotted path and the raw RHS string.
+type overrideEntry struct {
+	path string
+	raw  string
+}
+
+// validateOverridePaths checks every dotted path against the runConfig schema
+// before mutation, so an error names the exact --override key rather than the
+// re-decode's Go-type language.
+func validateOverridePaths(entries []overrideEntry) error {
+	for _, e := range entries {
+		if err := checkOverridePath(strings.Split(e.path, "."), reflect.TypeFor[runConfig](), e.path); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// checkOverridePath recursively validates one dotted path against a struct type
+// whose fields carry `yaml:` tags.
+func checkOverridePath(parts []string, t reflect.Type, fullPath string) error {
+	field := parts[0]
+	fields := yamlFields(t)
+	sub, ok := fields[field]
+	if !ok {
+		return fmt.Errorf("invalid --override %q: %q is not a known field; available fields are: %s",
+			fullPath, field, strings.Join(slices.Sorted(maps.Keys(fields)), ", "))
+	}
+	if len(parts) == 1 {
+		return nil
+	}
+	if freeFormFields[field] {
+		return nil
+	}
+	subStruct := underlyingStruct(sub)
+	if subStruct == nil {
+		return fmt.Errorf("invalid --override %q: %q is not a nested object; cannot address sub-field %q",
+			fullPath, field, strings.Join(parts[1:], "."))
+	}
+	return checkOverridePath(parts[1:], subStruct, fullPath)
+}
+
+// yamlFields maps a struct's yaml tag names to their field types, skipping
+// fields without a yaml tag (or tagged "-").
+func yamlFields(t reflect.Type) map[string]reflect.Type {
+	out := map[string]reflect.Type{}
+	for f := range t.Fields() {
+		tag := f.Tag.Get("yaml")
+		if tag == "" || tag == "-" {
+			continue
+		}
+		name, _, _ := strings.Cut(tag, ",")
+		if name == "" || name == "-" {
+			continue
+		}
+		out[name] = f.Type
+	}
+	return out
+}
+
+// underlyingStruct unwraps pointer/slice indirection and returns the struct type
+// a field decodes into, or nil if the field is not a struct (a scalar/map/etc.).
+func underlyingStruct(t reflect.Type) reflect.Type {
+	for t.Kind() == reflect.Pointer || t.Kind() == reflect.Slice {
+		t = t.Elem()
+	}
+	if t.Kind() == reflect.Struct {
+		return t
+	}
+	return nil
+}
+
+// applyOverrides walks each dotted path into the parsed YAML map and sets the
+// leaf to the RHS parsed as a YAML scalar. Intermediate maps are auto-created so
+// an override can add a field the YAML omits; the later re-decode rejects paths
+// absent from the schema. Changes are logged to stderr to keep JSON stdout clean.
+func applyOverrides(ctx context.Context, m map[string]any, entries []overrideEntry) error {
+	for _, e := range entries {
+		var value any
+		if err := yaml.Unmarshal([]byte(e.raw), &value); err != nil {
+			return fmt.Errorf("invalid --override %q: cannot parse value %q: %w", e.path, e.raw, err)
+		}
+
+		parts := strings.Split(e.path, ".")
+		current := m
+		for _, part := range parts[:len(parts)-1] {
+			next, ok := current[part].(map[string]any)
+			if !ok {
+				next = map[string]any{}
+				current[part] = next
+			}
+			current = next
+		}
+
+		leaf := parts[len(parts)-1]
+		old, had := current[leaf]
+		current[leaf] = value
+		if had {
+			logOverride(ctx, fmt.Sprintf("Override: changing %s from %v to %v", e.path, old, value))
+		} else {
+			logOverride(ctx, fmt.Sprintf("Override: setting %s to %v", e.path, value))
+		}
+	}
+	return nil
+}
+
+// logOverride writes to stderr only when a cmdIO is present; cmdio.LogString
+// panics without one, as in non-command callers such as unit tests.
+func logOverride(ctx context.Context, msg string) {
+	if cmdio.HasIO(ctx) {
+		cmdio.LogString(ctx, msg)
+	}
+}