Add template_dir support to template schema (#3686)

This allows template schemas to reference template files from a
different directory using the template_dir property. This enables schema
sharing between templates while keeping their template files separate.

## Changes

- Add TemplateDir field to jsonschema.Extension
- Update template readers to support template_dir resolution
- Modify template writer to use new LoadSchemaAndTemplateFS interface
- Add comprehensive tests for template_dir functionality

🤖 Extracted from #3671 with
[Claude Code](https://claude.ai/code)

## Tests
* Standard unit tests, acceptance tests in
#3671

---------

Co-authored-by: Claude <[email protected]>

Author: lennartkats-db

NEXT_CHANGELOG.md

@@ -9,6 +9,7 @@
 ### Dependency updates
 
 ### Bundles
+* Added support for a "template_dir" option in the databricks_template_schema.json format. ([#3671](https://github.com/databricks/cli/pull/3671)).
 * Remove resources.apps.config section ([#3680](https://github.com/databricks/cli/pull/3680))
 * Prompt for serverless compute in `dbt-sql` template (defaults to `yes`) ([#3668](https://github.com/databricks/cli/pull/3668))
 

libs/jsonschema/extension.go

@@ -35,6 +35,12 @@ type Extension struct {
 	// compatible with the current CLI version.
 	Version *int `json:"version,omitempty"`
 
+	// TemplateDir specifies the directory containing the template files to use.
+	// If not specified, the template files are expected to be in the same directory
+	// as the schema file. This allows schema files to reference template files
+	// from a different directory (e.g., "../default").
+	TemplateDir string `json:"template_dir,omitempty"`
+
 	// Preview indicates launch stage (e.g. PREVIEW).
 	//
 	// This field indicates whether the associated field is part of a private preview feature.

libs/template/config.go

@@ -34,6 +34,10 @@ func newConfig(ctx context.Context, templateFS fs.FS, schemaPath string) (*confi
 	if err != nil {
 		return nil, err
 	}
+	return newConfigFromSchema(ctx, schema)
+}
+
+func newConfigFromSchema(ctx context.Context, schema *jsonschema.Schema) (*config, error) {
 	if err := validateSchema(schema); err != nil {
 		return nil, err
 	}

libs/template/reader.go

@@ -10,40 +10,69 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/jsonschema"
 	"github.com/databricks/cli/libs/log"
 )
 
 type Reader interface {
-	// FS returns a file system that contains the template
-	// definition files.
-	FS(ctx context.Context) (fs.FS, error)
+	// LoadSchemaAndTemplateFS loads and returns the schema and template filesystem.
+	LoadSchemaAndTemplateFS(ctx context.Context) (*jsonschema.Schema, fs.FS, error)
 
 	// Cleanup releases any resources associated with the reader
 	// like cleaning up temporary directories.
 	Cleanup(ctx context.Context)
 }
 
+// builtinReader reads a template from the built-in templates.
 type builtinReader struct {
 	name string
 }
 
-func (r *builtinReader) FS(ctx context.Context) (fs.FS, error) {
+func (r *builtinReader) LoadSchemaAndTemplateFS(ctx context.Context) (*jsonschema.Schema, fs.FS, error) {
 	builtin, err := builtin()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
+	var schemaFS fs.FS
 	for _, entry := range builtin {
 		if entry.Name == r.name {
-			return entry.FS, nil
+			schemaFS = entry.FS
+			break
 		}
 	}
 
-	return nil, fmt.Errorf("builtin template %s not found", r.name)
+	if schemaFS == nil {
+		return nil, nil, fmt.Errorf("builtin template %s not found", r.name)
+	}
+
+	schema, err := jsonschema.LoadFS(schemaFS, schemaFileName)
+	if err != nil {
+		if errors.Is(err, fs.ErrNotExist) {
+			return nil, nil, fmt.Errorf("not a bundle template: expected to find a template schema file at %s", schemaFileName)
+		}
+		return nil, nil, fmt.Errorf("failed to load schema for template %s: %w", r.name, err)
+	}
+
+	// If no template_dir is specified, assume it's in the same directory as the schema
+	if schema.TemplateDir == "" {
+		return schema, schemaFS, nil
+	}
+
+	// Find the referenced template filesystem
+	templateDirName := filepath.Base(schema.TemplateDir)
+	for _, entry := range builtin {
+		if entry.Name == templateDirName {
+			return schema, entry.FS, nil
+		}
+	}
+
+	return nil, nil, fmt.Errorf("template directory %s (referenced by %s) not found", templateDirName, r.name)
 }
 
 func (r *builtinReader) Cleanup(ctx context.Context) {}
 
+// gitReader reads a template from a git repository.
 type gitReader struct {
 	gitUrl string
 	// tag or branch to checkout
@@ -66,19 +95,19 @@ func repoName(url string) string {
 	return parts[len(parts)-1]
 }
 
-func (r *gitReader) FS(ctx context.Context) (fs.FS, error) {
-	// Calling FS twice will lead to two downloaded copies of the git repo.
-	// In the future if you need to call FS twice, consider adding some caching
+func (r *gitReader) LoadSchemaAndTemplateFS(ctx context.Context) (*jsonschema.Schema, fs.FS, error) {
+	// Calling LoadSchemaAndTemplateFS twice will lead to two downloaded copies of the git repo.
+	// In the future if you need to call this twice, consider adding some caching
 	// logic here to avoid multiple downloads.
 	if r.tmpRepoDir != "" {
-		return nil, errors.New("FS called twice on git reader")
+		return nil, nil, errors.New("LoadSchemaAndTemplateFS called twice on git reader")
 	}
 
 	// Create a temporary directory with the name of the repository.  The '*'
 	// character is replaced by a random string in the generated temporary directory.
 	repoDir, err := os.MkdirTemp("", repoName(r.gitUrl)+"-*")
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	r.tmpRepoDir = repoDir
 
@@ -89,10 +118,11 @@ func (r *gitReader) FS(ctx context.Context) (fs.FS, error) {
 	err = r.cloneFunc(ctx, r.gitUrl, r.ref, repoDir)
 	close(promptSpinner)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	return os.DirFS(filepath.Join(repoDir, r.templateDir)), nil
+	templateDir := filepath.Join(repoDir, r.templateDir)
+	return loadSchemaAndResolveTemplateDir(templateDir)
 }
 
 func (r *gitReader) Cleanup(ctx context.Context) {
@@ -107,13 +137,42 @@ func (r *gitReader) Cleanup(ctx context.Context) {
 	}
 }
 
+// localReader reads a template from a local filesystem.
 type localReader struct {
 	// Path on the local filesystem that contains the template
 	path string
 }
 
-func (r *localReader) FS(ctx context.Context) (fs.FS, error) {
-	return os.DirFS(r.path), nil
+func (r *localReader) LoadSchemaAndTemplateFS(ctx context.Context) (*jsonschema.Schema, fs.FS, error) {
+	return loadSchemaAndResolveTemplateDir(r.path)
 }
 
 func (r *localReader) Cleanup(ctx context.Context) {}
+
+// loadSchemaAndResolveTemplateDir loads a schema from a local directory path
+// and resolves any template_dir reference.
+func loadSchemaAndResolveTemplateDir(path string) (*jsonschema.Schema, fs.FS, error) {
+	templateFS := os.DirFS(path)
+	schema, err := jsonschema.LoadFS(templateFS, schemaFileName)
+	if err != nil {
+		if errors.Is(err, fs.ErrNotExist) {
+			return nil, nil, fmt.Errorf("not a bundle template: expected to find a template schema file at %s", schemaFileName)
+		}
+		return nil, nil, fmt.Errorf("failed to load schema: %w", err)
+	}
+
+	// If no template_dir is specified, just use templateFS
+	if schema.TemplateDir == "" {
+		return schema, templateFS, nil
+	}
+
+	// Resolve template_dir relative to the schema location
+	templateDir := filepath.Join(path, schema.TemplateDir)
+
+	// Check if the referenced template directory exists
+	if _, err := os.Stat(templateDir); os.IsNotExist(err) {
+		return nil, nil, fmt.Errorf("template directory %s not found", templateDir)
+	}
+
+	return schema, os.DirFS(templateDir), nil
+}

libs/template/reader_test.go

@@ -23,25 +23,66 @@ func TestBuiltInReader(t *testing.T) {
 	for _, name := range exists {
 		t.Run(name, func(t *testing.T) {
 			r := &builtinReader{name: name}
-			fsys, err := r.FS(context.Background())
+			schema, fsys, err := r.LoadSchemaAndTemplateFS(context.Background())
 			assert.NoError(t, err)
 			assert.NotNil(t, fsys)
+			assert.NotNil(t, schema)
 
-			// Assert file content returned is accurate and every template has a welcome
-			// message defined.
-			b, err := fs.ReadFile(fsys, "databricks_template_schema.json")
-			require.NoError(t, err)
-			assert.Contains(t, string(b), "welcome_message")
+			// Assert schema has a welcome message defined.
+			assert.NotEmpty(t, schema.WelcomeMessage)
 		})
 	}
 
 	t.Run("doesnotexist", func(t *testing.T) {
 		r := &builtinReader{name: "doesnotexist"}
-		_, err := r.FS(context.Background())
+		_, _, err := r.LoadSchemaAndTemplateFS(context.Background())
 		assert.EqualError(t, err, "builtin template doesnotexist not found")
 	})
 }
 
+func TestBuiltInReaderTemplateDir(t *testing.T) {
+	// Test that template_dir property works correctly
+	// default-python template should use schema from default-python/ but template files from default/
+	r := &builtinReader{name: "default-python"}
+	schema, fsys, err := r.LoadSchemaAndTemplateFS(context.Background())
+	require.NoError(t, err)
+	assert.NotNil(t, schema)
+	assert.NotNil(t, fsys)
+
+	// Verify the schema contains default-python specific content
+	assert.Contains(t, schema.WelcomeMessage, "default Python template")
+
+	// Verify we can read template files (should come from default/)
+	templateFiles, err := fs.ReadDir(fsys, "template")
+	require.NoError(t, err)
+	assert.NotEmpty(t, templateFiles)
+
+	// Verify that a specific template file exists (this should come from default/ template)
+	_, err = fs.Stat(fsys, "template/{{.project_name}}/databricks.yml.tmpl")
+	assert.NoError(t, err)
+
+	// Test that a template without template_dir works normally
+	r2 := &builtinReader{name: "default-sql"}
+	schema2, fsys2, err := r2.LoadSchemaAndTemplateFS(context.Background())
+	require.NoError(t, err)
+	assert.NotNil(t, schema2)
+	assert.NotNil(t, fsys2)
+
+	// For default-sql, the schema should not reference template_dir
+	assert.Contains(t, schema2.WelcomeMessage, "default SQL template")
+
+	// Verify that lakeflow-pipelines also uses template_dir correctly
+	r3 := &builtinReader{name: "lakeflow-pipelines"}
+	schema3, fsys3, err := r3.LoadSchemaAndTemplateFS(context.Background())
+	require.NoError(t, err)
+	assert.NotNil(t, schema3)
+	assert.NotNil(t, fsys3)
+
+	// lakeflow-pipelines should also have template files from default/
+	_, err = fs.Stat(fsys3, "template/{{.project_name}}/databricks.yml.tmpl")
+	assert.NoError(t, err)
+}
+
 func TestGitUrlReader(t *testing.T) {
 	ctx := cmdio.MockDiscard(context.Background())
 
@@ -51,6 +92,7 @@ func TestGitUrlReader(t *testing.T) {
 		numCalls++
 		args = []string{url, reference, targetPath}
 		testutil.WriteFile(t, filepath.Join(targetPath, "a", "b", "c", "somefile"), "somecontent")
+		testutil.WriteFile(t, filepath.Join(targetPath, "a", "b", "c", "databricks_template_schema.json"), `{"welcome_message": "test"}`)
 		return nil
 	}
 	r := &gitReader{
@@ -61,21 +103,22 @@ func TestGitUrlReader(t *testing.T) {
 	}
 
 	// Assert cloneFunc is called with the correct args.
-	fsys, err := r.FS(ctx)
+	schema, fsys, err := r.LoadSchemaAndTemplateFS(ctx)
 	require.NoError(t, err)
 	require.NotEmpty(t, r.tmpRepoDir)
 	assert.Equal(t, 1, numCalls)
 	assert.DirExists(t, r.tmpRepoDir)
 	assert.Equal(t, []string{"someurl", "sometag", r.tmpRepoDir}, args)
+	assert.NotNil(t, schema)
 
 	// Assert the fs returned is rooted at the templateDir.
 	b, err := fs.ReadFile(fsys, "somefile")
 	require.NoError(t, err)
 	assert.Equal(t, "somecontent", string(b))
 
-	// Assert second call to FS returns an error.
-	_, err = r.FS(ctx)
-	assert.ErrorContains(t, err, "FS called twice on git reader")
+	// Assert second call returns an error.
+	_, _, err = r.LoadSchemaAndTemplateFS(ctx)
+	assert.ErrorContains(t, err, "LoadSchemaAndTemplateFS called twice on git reader")
 
 	// Assert the downloaded repository is cleaned up.
 	_, err = fs.Stat(fsys, ".")
@@ -88,14 +131,87 @@ func TestGitUrlReader(t *testing.T) {
 func TestLocalReader(t *testing.T) {
 	tmpDir := t.TempDir()
 	testutil.WriteFile(t, filepath.Join(tmpDir, "somefile"), "somecontent")
+	testutil.WriteFile(t, filepath.Join(tmpDir, "databricks_template_schema.json"), `{"welcome_message": "test"}`)
 	ctx := context.Background()
 
 	r := &localReader{path: tmpDir}
-	fsys, err := r.FS(ctx)
+	schema, fsys, err := r.LoadSchemaAndTemplateFS(ctx)
 	require.NoError(t, err)
+	assert.NotNil(t, schema)
 
 	// Assert the fs returned is rooted at correct location.
 	b, err := fs.ReadFile(fsys, "somefile")
 	require.NoError(t, err)
 	assert.Equal(t, "somecontent", string(b))
 }
+
+func TestLocalReaderWithTemplateDir(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	// Create a template directory with template_dir pointing to another directory
+	schemaDir := filepath.Join(tmpDir, "schema-template")
+	templateDir := filepath.Join(tmpDir, "actual-template")
+
+	// Create the schema template directory with a schema that references ../actual-template
+	testutil.WriteFile(t, filepath.Join(schemaDir, "databricks_template_schema.json"),
+		`{"welcome_message": "test with template_dir", "template_dir": "../actual-template"}`)
+
+	// Create the actual template directory with template files
+	testutil.WriteFile(t, filepath.Join(templateDir, "template", "somefile"), "content from template_dir")
+	testutil.WriteFile(t, filepath.Join(templateDir, "template", "{{.project_name}}", "test.yml.tmpl"), "test template content")
+
+	ctx := context.Background()
+	r := &localReader{path: schemaDir}
+	schema, fsys, err := r.LoadSchemaAndTemplateFS(ctx)
+	require.NoError(t, err)
+	assert.NotNil(t, schema)
+	assert.Equal(t, "test with template_dir", schema.WelcomeMessage)
+
+	// Assert the fs returned is rooted at the template_dir location
+	b, err := fs.ReadFile(fsys, "template/somefile")
+	require.NoError(t, err)
+	assert.Equal(t, "content from template_dir", string(b))
+
+	// Verify we can read the templated file
+	b2, err := fs.ReadFile(fsys, "template/{{.project_name}}/test.yml.tmpl")
+	require.NoError(t, err)
+	assert.Equal(t, "test template content", string(b2))
+}
+
+func TestGitReaderWithTemplateDir(t *testing.T) {
+	ctx := cmdio.MockDiscard(context.Background())
+
+	cloneFunc := func(ctx context.Context, url, reference, targetPath string) error {
+		// Create a template with template_dir reference
+		schemaDir := filepath.Join(targetPath, "a", "b", "c")
+		templateDir := filepath.Join(targetPath, "a", "b", "actual-template")
+
+		testutil.WriteFile(t, filepath.Join(schemaDir, "databricks_template_schema.json"),
+			`{"welcome_message": "git test with template_dir", "template_dir": "../actual-template"}`)
+
+		// Create the actual template directory with template files
+		testutil.WriteFile(t, filepath.Join(templateDir, "template", "gitfile"), "content from git template_dir")
+
+		return nil
+	}
+
+	r := &gitReader{
+		gitUrl:      "someurl",
+		cloneFunc:   cloneFunc,
+		ref:         "sometag",
+		templateDir: "a/b/c",
+	}
+
+	schema, fsys, err := r.LoadSchemaAndTemplateFS(ctx)
+	require.NoError(t, err)
+	assert.NotNil(t, schema)
+	assert.Equal(t, "git test with template_dir", schema.WelcomeMessage)
+
+	// Assert the fs returned is rooted at the template_dir location
+	b, err := fs.ReadFile(fsys, "template/gitfile")
+	require.NoError(t, err)
+	assert.Equal(t, "content from git template_dir", string(b))
+
+	// Cleanup
+	r.Cleanup(ctx)
+}