air: flatten code_source configuration

.nextchanges/air/code-source-flattening.md

@@ -0,0 +1 @@
+* Support snapshot fields directly under `code_source` without requiring `type` or a nested `snapshot` block. ([#6927](https://github.com/databricks/cli/pull/6927))

acceptance/air/config-help/output.txt

@@ -114,29 +114,34 @@ config.compute.accelerator_type
   Type:     string
   Required: no
 
-=== conditionally required leaf
->>> [CLI] air run -h config.code_source.type
-config.code_source.type
-  Kind of code source. Must be "snapshot", the only supported type.
+=== code source fields
+>>> [CLI] air run -h config.code_source
+config.code_source
+  Local code to upload and make available to the run.
 
-  Type:     string
-  Required: when code_source is set
+  Fields:
+    root_path      Root of the code source to archive.
+    remote_volume  Volume to upload the archive to.
+    git            Pin the snapshot to a specific git revision.
+    include_paths  Restrict the archive to these paths, relative to root_path and without "..".
 
-=== union field reports both accepted shapes
->>> [CLI] air run -h config.code_source.snapshot.git.remote
-config.code_source.snapshot.git.remote
+Use "-h config.code_source.<field>" for details on a field.
+
+=== union field
+>>> [CLI] air run -h config.code_source.git.remote
+config.code_source.git.remote
   No longer supported: the snapshot archives your local copy. Only false is accepted; use commit to pin a revision.
 
   Type:     bool or string
   Required: no
 
-=== git-pinned subdirectory snapshots document root_path scoping
->>> [CLI] air run -h config.code_source.snapshot.root_path
-config.code_source.snapshot.root_path
+=== git-pinned subdirectory documents root_path scoping
+>>> [CLI] air run -h config.code_source.root_path
+config.code_source.root_path
   Root of the code source to archive. A git-pinned subdirectory packages only that subtree.
 
   Type:     string
-  Required: when code_source.snapshot is set
+  Required: when code_source is set
 
 === the config. prefix is optional
 >>> [CLI] air run -h compute.num_accelerators

acceptance/air/config-help/script

@@ -22,14 +22,14 @@ trace $CLI air run -h config.environment.version
 title "leaf field"
 trace $CLI air run -h config.compute.accelerator_type
 
-title "conditionally required leaf"
-trace $CLI air run -h config.code_source.type
+title "code source fields"
+trace $CLI air run -h config.code_source
 
-title "union field reports both accepted shapes"
-trace $CLI air run -h config.code_source.snapshot.git.remote
+title "union field"
+trace $CLI air run -h config.code_source.git.remote
 
-title "git-pinned subdirectory snapshots document root_path scoping"
-trace $CLI air run -h config.code_source.snapshot.root_path
+title "git-pinned subdirectory documents root_path scoping"
+trace $CLI air run -h config.code_source.root_path
 
 title "the config. prefix is optional"
 trace $CLI air run -h compute.num_accelerators

cmd/air/convert_to_dabs.go

@@ -127,7 +127,7 @@ func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir st
 		// artifact location is set bundle-wide via workspace.artifact_path, not
 		// per-code-source, so a per-source Volume isn't representable here.
 		if snap.RemoteVolume != nil {
-			return nil, nil, errors.New("code_source.snapshot.remote_volume is not supported by convert-to-dabs; set workspace.artifact_path in the bundle instead")
+			return nil, nil, errors.New("remote_volume in code_source is not supported by convert-to-dabs; set workspace.artifact_path in the bundle instead")
 		}
 	}
 

cmd/air/convert_to_dabs_test.go

@@ -318,7 +318,7 @@ code_source:
 	loaded, err := loadRunConfig(path)
 	require.NoError(t, err)
 	_, _, err = convertToDabs(t.Context(), loaded, path, filepath.Dir(path))
-	require.ErrorContains(t, err, "remote_volume is not supported")
+	require.ErrorContains(t, err, "remote_volume in code_source is not supported")
 }
 
 // env_variables / secrets / parameters ride as sidecar files (the ai_runtime_task

cmd/air/runconfig.go

@@ -325,52 +325,72 @@ func (s *stringOrInt) UnmarshalYAML(node *yaml.Node) error {
 
 // codeSourceConfig is the `code_source` block. Only the "snapshot" type exists.
 type codeSourceConfig struct {
-	Type     string                `yaml:"type" help:"Kind of code source. Must be \"snapshot\", the only supported type." required:"when code_source is set"`
-	Snapshot *snapshotSourceConfig `yaml:"snapshot" help:"Which local directory to archive and upload." required:"when code_source.type is \"snapshot\""`
+	Type           *string               `yaml:"type,omitempty" hidden:"yes"`
+	Snapshot       *snapshotSourceConfig `yaml:",inline"`
+	LegacySnapshot *snapshotSourceConfig `yaml:"snapshot,omitempty" hidden:"yes"`
 }
 
 func (c *codeSourceConfig) validate() error {
-	if c.Type != "snapshot" {
-		return fmt.Errorf("code_source.type must be 'snapshot', got %q", c.Type)
+	if c.Type != nil && strings.TrimSpace(*c.Type) == "" {
+		return errors.New("code_source.type cannot be empty; omit it to use snapshot")
 	}
-	if c.Snapshot == nil {
-		return errors.New("code_source.type='snapshot' requires a snapshot configuration")
+	if c.Type != nil && *c.Type != "snapshot" {
+		return fmt.Errorf("code_source.type must be 'snapshot', got %q", *c.Type)
 	}
-	return c.Snapshot.validate()
+	if c.Snapshot != nil && c.LegacySnapshot != nil {
+		return errors.New("code_source cannot combine the nested 'snapshot' block with direct fields; remove the 'snapshot' wrapper and place all fields directly under code_source")
+	}
+
+	snapshot := c.Snapshot
+	prefix := "code_source"
+	if snapshot == nil {
+		snapshot = c.LegacySnapshot
+		prefix = "code_source.snapshot"
+	}
+	if snapshot == nil {
+		return errors.New("code_source requires snapshot configuration with root_path")
+	}
+	if err := snapshot.validate(prefix); err != nil {
+		return err
+	}
+
+	c.Snapshot = snapshot
+	c.LegacySnapshot = nil
+	return nil
 }
 
 // snapshotSourceConfig describes a local directory to tar and upload.
 type snapshotSourceConfig struct {
-	RootPath     string   `yaml:"root_path" help:"Root of the code source to archive. A git-pinned subdirectory packages only that subtree." required:"when code_source.snapshot is set"`
+	RootPath     string   `yaml:"root_path" help:"Root of the code source to archive. A git-pinned subdirectory packages only that subtree." required:"when code_source is set"`
 	RemoteVolume *string  `yaml:"remote_volume" help:"Volume to upload the archive to. Must start with /Volumes/."`
 	Git          *gitRef  `yaml:"git" help:"Pin the snapshot to a specific git revision."`
 	IncludePaths []string `yaml:"include_paths" help:"Restrict the archive to these paths, relative to root_path and without \"..\". Omit to include everything."`
 }
 
-func (s *snapshotSourceConfig) validate() error {
+func (s *snapshotSourceConfig) validate(prefix string) error {
 	if strings.TrimSpace(s.RootPath) == "" {
-		return errors.New("code_source.snapshot.root_path cannot be empty")
+		return fmt.Errorf("%s.root_path cannot be empty", prefix)
 	}
 
 	if s.RemoteVolume != nil && !strings.HasPrefix(*s.RemoteVolume, "/Volumes/") {
-		return errors.New("code_source.snapshot.remote_volume must start with '/Volumes/'")
+		return fmt.Errorf("%s.remote_volume must start with '/Volumes/'", prefix)
 	}
 
 	// A non-nil but empty include_paths is an explicit mistake (omit it instead).
 	if s.IncludePaths != nil && len(s.IncludePaths) == 0 {
-		return errors.New("code_source.snapshot.include_paths cannot be an empty list; either omit it or provide paths")
+		return fmt.Errorf("%s.include_paths cannot be an empty list; either omit it or provide paths", prefix)
 	}
 	for _, p := range s.IncludePaths {
 		p = strings.TrimSpace(p)
 		if p == "" {
-			return errors.New("code_source.snapshot.include_paths entry cannot be empty")
+			return fmt.Errorf("%s.include_paths entry cannot be empty", prefix)
 		}
 		if strings.HasPrefix(p, "/") {
-			return fmt.Errorf("code_source.snapshot.include_paths must be relative paths, got: %s", p)
+			return fmt.Errorf("%s.include_paths must be relative paths, got: %s", prefix, p)
 		}
 		// No parent traversal: snapshots must stay within root_path.
 		if slices.Contains(strings.Split(p, "/"), "..") {
-			return fmt.Errorf("code_source.snapshot.include_paths cannot contain '..' traversal, got: %s", p)
+			return fmt.Errorf("%s.include_paths cannot contain '..' traversal, got: %s", prefix, p)
 		}
 	}
 
@@ -520,6 +540,7 @@ type configField struct {
 	typeName string
 	required string
 	help     string
+	hidden   bool
 	// freeForm marks a user-keyed map (parameters/secrets/env_variables): no
 	// children, yet any sub-path into it is valid.
 	freeForm bool
@@ -555,7 +576,7 @@ func resolveConfigField(path string) (configField, error) {
 			return configField{}, fmt.Errorf("%q is not an object, so it has no field %q", current.path, part)
 		}
 		child, ok := findConfigChild(current.children, part)
-		if !ok {
+		if !ok || child.hidden {
 			return configField{}, unknownConfigFieldError(current, part, strings.Split(trimmed, ".")[:i+1])
 		}
 		current = child
@@ -596,6 +617,9 @@ func cutLast(s, sep string) (before, after string, found bool) {
 func unknownConfigFieldError(parent configField, part string, matched []string) error {
 	names := make([]string, 0, len(parent.children))
 	for _, c := range parent.children {
+		if c.hidden {
+			continue
+		}
 		names = append(names, configLeafName(c.path))
 	}
 	slices.Sort(names)
@@ -653,7 +677,13 @@ func describeStruct(t reflect.Type, prefix string) []configField {
 		if tag == "" || tag == "-" {
 			continue
 		}
-		name, _, _ := strings.Cut(tag, ",")
+		name, options, _ := strings.Cut(tag, ",")
+		if slices.Contains(strings.Split(options, ","), "inline") {
+			if nested := underlyingConfigStruct(f.Type); nested != nil {
+				out = append(out, describeStruct(nested, prefix)...)
+			}
+			continue
+		}
 		if name == "" || name == "-" {
 			continue
 		}
@@ -663,6 +693,7 @@ func describeStruct(t reflect.Type, prefix string) []configField {
 			typeName: configTypeName(f.Type),
 			required: f.Tag.Get("required"),
 			help:     f.Tag.Get("help"),
+			hidden:   f.Tag.Get("hidden") == "yes",
 			freeForm: freeFormConfigFields[name],
 		}
 		if nested := underlyingConfigStruct(f.Type); nested != nil && !field.freeForm {
@@ -737,10 +768,16 @@ func renderConfigField(w io.Writer, f configField) {
 
 	width := 0
 	for _, c := range f.children {
+		if c.hidden {
+			continue
+		}
 		width = max(width, len(configLeafName(c.path)))
 	}
 	fmt.Fprintf(w, "\n  Fields:\n")
 	for _, c := range f.children {
+		if c.hidden {
+			continue
+		}
 		fmt.Fprintf(w, "    %-*s  %s\n", width, configLeafName(c.path), configFieldSummary(c))
 	}
 	fmt.Fprintf(w, "\nUse \"-h %s.<field>\" for details on a field.\n", f.path)

cmd/air/runconfig_load.go

@@ -102,6 +102,15 @@ func finishRunConfigLoad(path string, document *yaml.Node, validationYAML []byte
 	if err != nil {
 		return nil, err
 	}
+	// YAML decodes omitted and null pointers alike. Preserve explicit nulls so validation rejects them.
+	if cfg.CodeSource != nil {
+		if cfg.CodeSource.Type == nil && codeSourceFieldPresent(document, "type") {
+			cfg.CodeSource.Type = new(string)
+		}
+		if cfg.CodeSource.LegacySnapshot == nil && codeSourceFieldPresent(document, "snapshot") {
+			cfg.CodeSource.LegacySnapshot = &snapshotSourceConfig{}
+		}
+	}
 	if err := validateRunConfig(cfg); err != nil {
 		return nil, err
 	}
@@ -115,6 +124,15 @@ func finishRunConfigLoad(path string, document *yaml.Node, validationYAML []byte
 	return cfg, nil
 }
 
+func codeSourceFieldPresent(document *yaml.Node, field string) bool {
+	codeSource, _ := mappingValue(document.Content[0], "code_source")
+	if codeSource == nil || codeSource.Kind != yaml.MappingNode {
+		return false
+	}
+	fieldNode, _ := mappingValue(codeSource, field)
+	return fieldNode != nil
+}
+
 func stripYAMLComments(node *yaml.Node) {
 	node.HeadComment = ""
 	node.LineComment = ""

cmd/air/runconfig_override.go

@@ -67,6 +67,9 @@ func checkOverridePath(parts []string, node configField, fullPath string) error
 	if !ok {
 		names := make([]string, 0, len(node.children))
 		for _, c := range node.children {
+			if c.hidden {
+				continue
+			}
 			names = append(names, configLeafName(c.path))
 		}
 		slices.Sort(names)

cmd/air/runconfig_override_test.go

@@ -138,13 +138,15 @@ func TestLoadRunConfigWithOverrides(t *testing.T) {
 		assert.Equal(t, "0", cfg.EnvVariables["RANK"])
 	})
 
-	t.Run("intermediate maps are auto-created", func(t *testing.T) {
-		cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"code_source.type=snapshot", "code_source.snapshot.root_path=/Workspace/x"})
-		require.NoError(t, err)
-		require.NotNil(t, cfg.CodeSource)
-		require.NotNil(t, cfg.CodeSource.Snapshot)
-		assert.Equal(t, "/Workspace/x", cfg.CodeSource.Snapshot.RootPath)
-	})
+	for _, path := range []string{"code_source.root_path", "code_source.snapshot.root_path"} {
+		t.Run(path+" override", func(t *testing.T) {
+			cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{path + "=/Workspace/x"})
+			require.NoError(t, err)
+			require.NotNil(t, cfg.CodeSource)
+			require.NotNil(t, cfg.CodeSource.Snapshot)
+			assert.Equal(t, "/Workspace/x", cfg.CodeSource.Snapshot.RootPath)
+		})
+	}
 
 	t.Run("unity catalog image override applies", func(t *testing.T) {
 		cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"environment.unity_catalog_image=main.air.training:prod"})