From dd1d540429150a848425835df615ac38cc4914b2 Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Tue, 28 Nov 2023 10:04:06 +0100
Subject: [PATCH] Add mlops-stacks to the default `databricks bundle init`
 prompt (#988)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Changes
This makes mlops-stacks more discoverable and makes the UX of
initialising the mlops-stack template better.

## Tests
Manually

Dropdown UI:
```
shreyas.goenka@THW32HFW6T projects % cli bundle init
Template to use:
  ▸ default-python
    mlops-stacks
```

Help message:
```
shreyas.goenka@THW32HFW6T bricks % cli bundle init -h
Initialize using a bundle template.

TEMPLATE_PATH optionally specifies which template to use. It can be one of the following:
- default-python: The default Python template
- mlops-stacks: The Databricks MLOps Stacks template. More information can be found at: https://github.com/databricks/mlops-stacks
```
---
 cmd/bundle/init.go      | 66 ++++++++++++++++++++++++++++++++++-------
 cmd/bundle/init_test.go | 15 ++++++++++
 2 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/cmd/bundle/init.go b/cmd/bundle/init.go
index ccdc68591..ac6f49de3 100644
--- a/cmd/bundle/init.go
+++ b/cmd/bundle/init.go
@@ -2,8 +2,10 @@ package bundle
 
 import (
 	"errors"
+	"fmt"
 	"os"
 	"path/filepath"
+	"slices"
 	"strings"
 
 	"github.com/databricks/cli/cmd/root"
@@ -18,9 +20,52 @@ var gitUrlPrefixes = []string{
 	"git@",
 }
 
-var aliasedTemplates = map[string]string{
-	"mlops-stack":  "https://github.com/databricks/mlops-stacks",
-	"mlops-stacks": "https://github.com/databricks/mlops-stacks",
+type nativeTemplate struct {
+	name        string
+	gitUrl      string
+	description string
+	aliases     []string
+}
+
+var nativeTemplates = []nativeTemplate{
+	{
+		name:        "default-python",
+		description: "The default Python template",
+	},
+	{
+		name:        "mlops-stacks",
+		gitUrl:      "https://github.com/databricks/mlops-stacks",
+		description: "The Databricks MLOps Stacks template (https://github.com/databricks/mlops-stacks)",
+		aliases:     []string{"mlops-stack"},
+	},
+}
+
+func nativeTemplateDescriptions() string {
+	var lines []string
+	for _, template := range nativeTemplates {
+		lines = append(lines, fmt.Sprintf("- %s: %s", template.name, template.description))
+	}
+	return strings.Join(lines, "\n")
+}
+
+func nativeTemplateOptions() []string {
+	names := make([]string, 0, len(nativeTemplates))
+	for _, template := range nativeTemplates {
+		names = append(names, template.name)
+	}
+	return names
+}
+
+func getUrlForNativeTemplate(name string) string {
+	for _, template := range nativeTemplates {
+		if template.name == name {
+			return template.gitUrl
+		}
+		if slices.Contains(template.aliases, name) {
+			return template.gitUrl
+		}
+	}
+	return ""
 }
 
 func isRepoUrl(url string) bool {
@@ -47,14 +92,14 @@ func newInitCommand() *cobra.Command {
 		Use:   "init [TEMPLATE_PATH]",
 		Short: "Initialize using a bundle template",
 		Args:  cobra.MaximumNArgs(1),
-		Long: `Initialize using a bundle template.
+		Long: fmt.Sprintf(`Initialize using a bundle template.
 
 TEMPLATE_PATH optionally specifies which template to use. It can be one of the following:
-- 'default-python' for the default Python template
+%s
 - a local file system path with a template directory
 - a Git repository URL, e.g. https://github.com/my/repository
 
-See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more information on templates.`,
+See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more information on templates.`, nativeTemplateDescriptions()),
 	}
 
 	var configFile string
@@ -89,15 +134,16 @@ See https://docs.databricks.com/en/dev-tools/bundles/templates.html for more inf
 			if !cmdio.IsOutTTY(ctx) || !cmdio.IsInTTY(ctx) {
 				return errors.New("please specify a template")
 			}
-			templatePath, err = cmdio.Ask(ctx, "Template to use", "default-python")
+			templatePath, err = cmdio.AskSelect(ctx, "Template to use", nativeTemplateOptions())
 			if err != nil {
 				return err
 			}
 		}
 
-		// Expand templatePath if it's an alias for a known template
-		if _, ok := aliasedTemplates[templatePath]; ok {
-			templatePath = aliasedTemplates[templatePath]
+		// Expand templatePath to a git URL if it's an alias for a known native template
+		// and we know it's git URL.
+		if gitUrl := getUrlForNativeTemplate(templatePath); gitUrl != "" {
+			templatePath = gitUrl
 		}
 
 		if !isRepoUrl(templatePath) {
diff --git a/cmd/bundle/init_test.go b/cmd/bundle/init_test.go
index 4a795160e..db4446bb1 100644
--- a/cmd/bundle/init_test.go
+++ b/cmd/bundle/init_test.go
@@ -25,3 +25,18 @@ func TestBundleInitRepoName(t *testing.T) {
 	assert.Equal(t, "invalid-url", repoName("invalid-url"))
 	assert.Equal(t, "www.github.com", repoName("https://www.github.com"))
 }
+
+func TestNativeTemplateOptions(t *testing.T) {
+	assert.Equal(t, []string{"default-python", "mlops-stacks"}, nativeTemplateOptions())
+}
+
+func TestNativeTemplateDescriptions(t *testing.T) {
+	assert.Equal(t, "- default-python: The default Python template\n- mlops-stacks: The Databricks MLOps Stacks template (https://github.com/databricks/mlops-stacks)", nativeTemplateDescriptions())
+}
+
+func TestGetUrlForNativeTemplate(t *testing.T) {
+	assert.Equal(t, "https://github.com/databricks/mlops-stacks", getUrlForNativeTemplate("mlops-stacks"))
+	assert.Equal(t, "https://github.com/databricks/mlops-stacks", getUrlForNativeTemplate("mlops-stack"))
+	assert.Equal(t, "", getUrlForNativeTemplate("default-python"))
+	assert.Equal(t, "", getUrlForNativeTemplate("invalid"))
+}