From e2174f5870af423e055dfd59d5967da5f553390c Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Mon, 18 Nov 2024 20:35:36 +0100
Subject: [PATCH] Make init command aware of built-in templates

---
 cmd/bundle/init.go            | 35 ++++++++++++++++++++++++--
 libs/template/builtin.go      | 47 +++++++++++++++++++++++++++++++++++
 libs/template/builtin_test.go | 28 +++++++++++++++++++++
 libs/template/materialize.go  |  4 ---
 4 files changed, 108 insertions(+), 6 deletions(-)
 create mode 100644 libs/template/builtin.go
 create mode 100644 libs/template/builtin_test.go

diff --git a/cmd/bundle/init.go b/cmd/bundle/init.go
index 7f2c0efc..d31a702a 100644
--- a/cmd/bundle/init.go
+++ b/cmd/bundle/init.go
@@ -3,6 +3,7 @@ package bundle
 import (
 	"errors"
 	"fmt"
+	"io/fs"
 	"os"
 	"path/filepath"
 	"slices"
@@ -109,6 +110,24 @@ func getUrlForNativeTemplate(name string) string {
 	return ""
 }
 
+func getFsForNativeTemplate(name string) (fs.FS, error) {
+	builtin, err := template.Builtin()
+	if err != nil {
+		return nil, err
+	}
+
+	// If this is a built-in template, the return value will be non-nil.
+	var templateFS fs.FS
+	for _, entry := range builtin {
+		if entry.Name == name {
+			templateFS = entry.FS
+			break
+		}
+	}
+
+	return templateFS, nil
+}
+
 func isRepoUrl(url string) bool {
 	result := false
 	for _, prefix := range gitUrlPrefixes {
@@ -198,9 +217,20 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 			if templateDir != "" {
 				return errors.New("--template-dir can only be used with a Git repository URL")
 			}
+
+			templateFS, err := getFsForNativeTemplate(templatePath)
+			if err != nil {
+				return err
+			}
+
+			// If this is not a built-in template, then it must be a local file system path.
+			if templateFS == nil {
+				templateFS = os.DirFS(templatePath)
+			}
+
 			// skip downloading the repo because input arg is not a URL. We assume
 			// it's a path on the local file system in that case
-			return template.Materialize(ctx, configFile, templatePath, outputDir)
+			return template.Materialize(ctx, configFile, templateFS, outputDir)
 		}
 
 		// Create a temporary directory with the name of the repository.  The '*'
@@ -224,7 +254,8 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 
 		// Clean up downloaded repository once the template is materialized.
 		defer os.RemoveAll(repoDir)
-		return template.Materialize(ctx, configFile, filepath.Join(repoDir, templateDir), outputDir)
+		templateFS := os.DirFS(filepath.Join(repoDir, templateDir))
+		return template.Materialize(ctx, configFile, templateFS, outputDir)
 	}
 	return cmd
 }
diff --git a/libs/template/builtin.go b/libs/template/builtin.go
new file mode 100644
index 00000000..dcb3a885
--- /dev/null
+++ b/libs/template/builtin.go
@@ -0,0 +1,47 @@
+package template
+
+import (
+	"embed"
+	"io/fs"
+)
+
+//go:embed all:templates
+var builtinTemplates embed.FS
+
+// BuiltinTemplate represents a template that is built into the CLI.
+type BuiltinTemplate struct {
+	Name string
+	FS   fs.FS
+}
+
+// Builtin returns the list of all built-in templates.
+func Builtin() ([]BuiltinTemplate, error) {
+	templates, err := fs.Sub(builtinTemplates, "templates")
+	if err != nil {
+		return nil, err
+	}
+
+	entries, err := fs.ReadDir(templates, ".")
+	if err != nil {
+		return nil, err
+	}
+
+	var out []BuiltinTemplate
+	for _, entry := range entries {
+		if !entry.IsDir() {
+			continue
+		}
+
+		templateFS, err := fs.Sub(templates, entry.Name())
+		if err != nil {
+			return nil, err
+		}
+
+		out = append(out, BuiltinTemplate{
+			Name: entry.Name(),
+			FS:   templateFS,
+		})
+	}
+
+	return out, nil
+}
diff --git a/libs/template/builtin_test.go b/libs/template/builtin_test.go
new file mode 100644
index 00000000..504e0acc
--- /dev/null
+++ b/libs/template/builtin_test.go
@@ -0,0 +1,28 @@
+package template
+
+import (
+	"io/fs"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBuiltin(t *testing.T) {
+	out, err := Builtin()
+	require.NoError(t, err)
+	assert.Len(t, out, 3)
+
+	// Confirm names.
+	assert.Equal(t, "dbt-sql", out[0].Name)
+	assert.Equal(t, "default-python", out[1].Name)
+	assert.Equal(t, "default-sql", out[2].Name)
+
+	// Confirm that the filesystems work.
+	_, err = fs.Stat(out[0].FS, `template/{{.project_name}}/dbt_project.yml.tmpl`)
+	assert.NoError(t, err)
+	_, err = fs.Stat(out[1].FS, `template/{{.project_name}}/tests/main_test.py.tmpl`)
+	assert.NoError(t, err)
+	_, err = fs.Stat(out[2].FS, `template/{{.project_name}}/src/orders_daily.sql.tmpl`)
+	assert.NoError(t, err)
+}
diff --git a/libs/template/materialize.go b/libs/template/materialize.go
index e26a599b..c24519ef 100644
--- a/libs/template/materialize.go
+++ b/libs/template/materialize.go
@@ -2,7 +2,6 @@ package template
 
 import (
 	"context"
-	"embed"
 	"errors"
 	"fmt"
 	"io/fs"
@@ -14,9 +13,6 @@ const libraryDirName = "library"
 const templateDirName = "template"
 const schemaFileName = "databricks_template_schema.json"
 
-//go:embed all:templates
-var builtinTemplates embed.FS
-
 // This function materializes the input templates as a project, using user defined
 // configurations.
 // Parameters: