Add mlops-stacks to the default `databricks bundle init` prompt (#988)

## 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
```
This commit is contained in:
shreyas-goenka 2023-11-28 10:04:06 +01:00 committed by GitHub
parent 96e9545cf0
commit dd1d540429
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 71 additions and 10 deletions

View File

@ -2,8 +2,10 @@ package bundle
import ( import (
"errors" "errors"
"fmt"
"os" "os"
"path/filepath" "path/filepath"
"slices"
"strings" "strings"
"github.com/databricks/cli/cmd/root" "github.com/databricks/cli/cmd/root"
@ -18,9 +20,52 @@ var gitUrlPrefixes = []string{
"git@", "git@",
} }
var aliasedTemplates = map[string]string{ type nativeTemplate struct {
"mlops-stack": "https://github.com/databricks/mlops-stacks", name string
"mlops-stacks": "https://github.com/databricks/mlops-stacks", 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 { func isRepoUrl(url string) bool {
@ -47,14 +92,14 @@ func newInitCommand() *cobra.Command {
Use: "init [TEMPLATE_PATH]", Use: "init [TEMPLATE_PATH]",
Short: "Initialize using a bundle template", Short: "Initialize using a bundle template",
Args: cobra.MaximumNArgs(1), 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: 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 local file system path with a template directory
- a Git repository URL, e.g. https://github.com/my/repository - 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 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) { if !cmdio.IsOutTTY(ctx) || !cmdio.IsInTTY(ctx) {
return errors.New("please specify a template") 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 { if err != nil {
return err return err
} }
} }
// Expand templatePath if it's an alias for a known template // Expand templatePath to a git URL if it's an alias for a known native template
if _, ok := aliasedTemplates[templatePath]; ok { // and we know it's git URL.
templatePath = aliasedTemplates[templatePath] if gitUrl := getUrlForNativeTemplate(templatePath); gitUrl != "" {
templatePath = gitUrl
} }
if !isRepoUrl(templatePath) { if !isRepoUrl(templatePath) {

View File

@ -25,3 +25,18 @@ func TestBundleInitRepoName(t *testing.T) {
assert.Equal(t, "invalid-url", repoName("invalid-url")) assert.Equal(t, "invalid-url", repoName("invalid-url"))
assert.Equal(t, "www.github.com", repoName("https://www.github.com")) 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"))
}