Add a foundation for built-in templates

.gitignore

@@ -28,3 +28,6 @@ __pycache__
 .terraform.lock.hcl
 
 .vscode/launch.json
+.vscode/tasks.json
+
+.databricks

.vscode/__builtins__.pyi

@@ -0,0 +1,3 @@
+# Typings for Pylance in VS Code
+# see https://github.com/microsoft/pyright/blob/main/docs/builtins.md
+from databricks.sdk.runtime import *

.vscode/settings.json

@@ -7,5 +7,6 @@
  "files.insertFinalNewline": true,
  "files.trimFinalNewlines": true,
  "python.envFile": "${workspaceFolder}/.databricks/.databricks.env",
- "databricks.python.envFile": "${workspaceFolder}/.env"
+ "databricks.python.envFile": "${workspaceFolder}/.env",
+ "python.analysis.stubPath": ".vscode"
 }

bundle/config/mutator/process_target_mode.go

@@ -8,7 +8,8 @@ import (
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
-	"github.com/databricks/databricks-sdk-go/service/iam"
+	"github.com/databricks/cli/libs/auth"
+	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/databricks/databricks-sdk-go/service/ml"
 )
@@ -111,7 +112,7 @@ func findIncorrectPath(b *bundle.Bundle, mode config.Mode) string {
 func validateProductionMode(ctx context.Context, b *bundle.Bundle, isPrincipalUsed bool) error {
 	if b.Config.Bundle.Git.Inferred {
  env := b.Config.Bundle.Target
-return fmt.Errorf("target with 'mode: production' must specify an explicit 'targets.%s.git' configuration", env)
+log.Warnf(ctx, "target with 'mode: production' should specify an explicit 'targets.%s.git' configuration", env)
 	}
 
 	r := b.Config.Resources
@@ -138,21 +139,6 @@ func validateProductionMode(ctx context.Context, b *bundle.Bundle, isPrincipalUs
 	return nil
 }
 
-// Determines whether a service principal identity is used to run the CLI.
-func isServicePrincipalUsed(ctx context.Context, b *bundle.Bundle) (bool, error) {
-	ws := b.WorkspaceClient()
-
-	// Check if a principal with the current user's ID exists.
-	// We need to use the ListAll method since Get is only usable by admins.
-	matches, err := ws.ServicePrincipals.ListAll(ctx, iam.ListServicePrincipalsRequest{
- Filter: "id eq " + b.Config.Workspace.CurrentUser.Id,
-	})
-	if err != nil {
- return false, err
-	}
-	return len(matches) > 0, nil
-}
-
 // Determines whether run_as is explicitly set for all resources.
 // We do this in a best-effort fashion rather than check the top-level
 // 'run_as' field because the latter is not required to be set.
@@ -174,10 +160,7 @@ func (m *processTargetMode) Apply(ctx context.Context, b *bundle.Bundle) error {
  }
  return transformDevelopmentMode(b)
 	case config.Production:
- isPrincipal, err := isServicePrincipalUsed(ctx, b)
- if err != nil {
- return err
- }
+ isPrincipal := auth.IsServicePrincipal(ctx, b.WorkspaceClient(), b.Config.Workspace.CurrentUser.Id)
  return validateProductionMode(ctx, b, isPrincipal)
 	case "":
  // No action

bundle/config/mutator/process_target_mode_test.go

@@ -118,17 +118,6 @@ func TestProcessTargetModeProduction(t *testing.T) {
 	assert.False(t, bundle.Config.Resources.Pipelines["pipeline1"].PipelineSpec.Development)
 }
 
-func TestProcessTargetModeProductionGit(t *testing.T) {
-	bundle := mockBundle(config.Production)
-
-	// Pretend the user didn't set Git configuration explicitly
-	bundle.Config.Bundle.Git.Inferred = true
-
-	err := validateProductionMode(context.Background(), bundle, false)
-	require.ErrorContains(t, err, "git")
-	bundle.Config.Bundle.Git.Inferred = false
-}
-
 func TestProcessTargetModeProductionOkForPrincipal(t *testing.T) {
 	bundle := mockBundle(config.Production)
 

cmd/bundle/init.go

@@ -1,10 +1,12 @@
 package bundle
 
 import (
+	"errors"
 	"os"
 	"path/filepath"
 	"strings"
 
+	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/git"
 	"github.com/databricks/cli/libs/template"
 	"github.com/spf13/cobra"
@@ -36,9 +38,9 @@ func repoName(url string) string {
 
 func newInitCommand() *cobra.Command {
 	cmd := &cobra.Command{
- Use: "init TEMPLATE_PATH",
+ Use: "init [TEMPLATE_PATH]",
  Short: "Initialize Template",
- Args: cobra.ExactArgs(1),
+ Args: cobra.MaximumNArgs(1),
 	}
 
 	var configFile string
@@ -48,9 +50,26 @@ func newInitCommand() *cobra.Command {
 	cmd.Flags().StringVar(&templateDir, "template-dir", "", "Directory within repository that holds the template specification.")
 	cmd.Flags().StringVar(&outputDir, "output-dir", "", "Directory to write the initialized template to.")
 
+	cmd.PreRunE = root.MustWorkspaceClient
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
- templatePath := args[0]
  ctx := cmd.Context()
+ var templatePath string
+ if len(args) > 0 {
+ templatePath = args[0]
+ } else {
+ return errors.New("please specify a template")
+
+ /* TODO: propose to use default-python (once template is ready)
+ var err error
+ if !cmdio.IsOutTTY(ctx) || !cmdio.IsInTTY(ctx) {
+ return errors.New("please specify a template")
+ }
+ templatePath, err = cmdio.Ask(ctx, "Template to use", "default-python")
+ if err != nil {
+ return err
+ }
+ */
+ }
 
  if !isRepoUrl(templatePath) {
  // skip downloading the repo because input arg is not a URL. We assume

cmd/root/auth.go

@@ -113,6 +113,10 @@ TRY_AUTH: // or try picking a config profile dynamically
 	return nil
 }
 
+func SetWorkspaceClient(ctx context.Context, w *databricks.WorkspaceClient) context.Context {
+	return context.WithValue(ctx, &workspaceClient, w)
+}
+
 func transformLoadError(path string, err error) error {
 	if os.IsNotExist(err) {
  return fmt.Errorf("no configuration file found at %s; please create one first", path)

cmd/root/bundle.go

@@ -43,7 +43,7 @@ func getTarget(cmd *cobra.Command) (value string) {
 	return target
 }
 
-func getProfile(cmd *cobra.Command) (value string) {
+func GetProfile(cmd *cobra.Command) (value string) {
 	// The command line flag takes precedence.
 	flag := cmd.Flag("profile")
 	if flag != nil {
@@ -70,7 +70,7 @@ func loadBundle(cmd *cobra.Command, args []string, load func(ctx context.Context
  return nil, nil
 	}
 
-	profile := getProfile(cmd)
+	profile := GetProfile(cmd)
 	if profile != "" {
  b.Config.Workspace.Profile = profile
 	}

libs/auth/service_principal.go

@@ -0,0 +1,16 @@
+package auth
+
+import (
+	"context"
+
+	"github.com/databricks/databricks-sdk-go"
+)
+
+// Determines whether a given user id is a service principal.
+// This function uses a heuristic: if no user exists with this id, we assume
+// it's a service principal. Unfortunately, the standard service principal API is too
+// slow for our purposes.
+func IsServicePrincipal(ctx context.Context, ws *databricks.WorkspaceClient, userId string) bool {
+	_, err := ws.Users.GetById(ctx, userId)
+	return err != nil
+}

libs/template/helpers.go

@@ -1,10 +1,16 @@
 package template
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"net/url"
 	"regexp"
 	"text/template"
+
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/auth"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 )
 
 type ErrFail struct {
@@ -20,35 +26,87 @@ type pair struct {
 	v any
 }
 
-var helperFuncs = template.FuncMap{
-	"fail": func(format string, args ...any) (any, error) {
- return nil, ErrFail{fmt.Sprintf(format, args...)}
-	},
-	// Alias for https://pkg.go.dev/net/url#Parse. Allows usage of all methods of url.URL
-	"url": func(rawUrl string) (*url.URL, error) {
- return url.Parse(rawUrl)
-	},
-	// Alias for https://pkg.go.dev/regexp#Compile. Allows usage of all methods of regexp.Regexp
-	"regexp": func(expr string) (*regexp.Regexp, error) {
- return regexp.Compile(expr)
-	},
-	// A key value pair. This is used with the map function to generate maps
-	// to use inside a template
-	"pair": func(k string, v any) pair {
- return pair{k, v}
-	},
-	// map converts a list of pairs to a map object. This is useful to pass multiple
-	// objects to templates defined in the library directory. Go text template
-	// syntax for invoking a template only allows specifying a single argument,
-	// this function can be used to workaround that limitation.
-	//
-	// For example: {{template "my_template" (map (pair "foo" $arg1) (pair "bar" $arg2))}}
-	// $arg1 and $arg2 can be referred from inside "my_template" as ".foo" and ".bar"
-	"map": func(pairs ...pair) map[string]any {
- result := make(map[string]any, 0)
- for _, p := range pairs {
- result[p.k] = p.v
- }
- return result
-	},
+func loadHelpers(ctx context.Context) template.FuncMap {
+	var user *iam.User
+	var is_service_principal *bool
+	w := root.WorkspaceClient(ctx)
+	return template.FuncMap{
+ "fail": func(format string, args ...any) (any, error) {
+ return nil, ErrFail{fmt.Sprintf(format, args...)}
+ },
+ // Alias for https://pkg.go.dev/net/url#Parse. Allows usage of all methods of url.URL
+ "url": func(rawUrl string) (*url.URL, error) {
+ return url.Parse(rawUrl)
+ },
+ // Alias for https://pkg.go.dev/regexp#Compile. Allows usage of all methods of regexp.Regexp
+ "regexp": func(expr string) (*regexp.Regexp, error) {
+ return regexp.Compile(expr)
+ },
+ // A key value pair. This is used with the map function to generate maps
+ // to use inside a template
+ "pair": func(k string, v any) pair {
+ return pair{k, v}
+ },
+ // map converts a list of pairs to a map object. This is useful to pass multiple
+ // objects to templates defined in the library directory. Go text template
+ // syntax for invoking a template only allows specifying a single argument,
+ // this function can be used to workaround that limitation.
+ //
+ // For example: {{template "my_template" (map (pair "foo" $arg1) (pair "bar" $arg2))}}
+ // $arg1 and $arg2 can be referred from inside "my_template" as ".foo" and ".bar"
+ "map": func(pairs ...pair) map[string]any {
+ result := make(map[string]any, 0)
+ for _, p := range pairs {
+ result[p.k] = p.v
+ }
+ return result
+ },
+ // Get smallest node type (follows Terraform's GetSmallestNodeType)
+ "smallest_node_type": func() (string, error) {
+ if w.Config.Host == "" {
+ return "", errors.New("cannot determine target workspace, please first setup a configuration profile using 'databricks auth login'")
+ }
+ if w.Config.IsAzure() {
+ return "Standard_D3_v2", nil
+ } else if w.Config.IsGcp() {
+ return "n1-standard-4", nil
+ }
+ return "i3.xlarge", nil
+ },
+ "workspace_host": func() (string, error) {
+ if w.Config.Host == "" {
+ return "", errors.New("cannot determine target workspace, please first setup a configuration profile using 'databricks auth login'")
+ }
+ return w.Config.Host, nil
+ },
+ "user_name": func() (string, error) {
+ if user == nil {
+ var err error
+ user, err = w.CurrentUser.Me(ctx)
+ if err != nil {
+ return "", err
+ }
+ }
+ result := user.UserName
+ if result == "" {
+ result = user.Id
+ }
+ return result, nil
+ },
+ "is_service_principal": func() (bool, error) {
+ if is_service_principal != nil {
+ return *is_service_principal, nil
+ }
+ if user == nil {
+ var err error
+ user, err = w.CurrentUser.Me(ctx)
+ if err != nil {
+ return false, err
+ }
+ }
+ result := auth.IsServicePrincipal(ctx, w, user.Id)
+ is_service_principal = &result
+ return result, nil
+ },
+	}
 }