properly pass along databricks profile

Author: fjakobs

experimental/apps-mcp/lib/middlewares/databricks_client.go

@@ -9,13 +9,15 @@ import (
 	"github.com/databricks/cli/experimental/apps-mcp/lib/mcp"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/prompts"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/session"
+	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/httpclient"
 )
 
 const (
-	DatabricksClientKey = "databricks_client"
+	DatabricksClientKey  = "databricks_client"
+	DatabricksProfileKey = "databricks_profile"
 )
 
 func NewDatabricksClientMiddleware(unauthorizedToolNames []string) mcp.Middleware {
@@ -40,8 +42,41 @@ func NewDatabricksClientMiddleware(unauthorizedToolNames []string) mcp.Middlewar
 	})
 }
 
-func MustGetApiClient(ctx context.Context) (*httpclient.ApiClient, error) {
-	w := MustGetDatabricksClient(ctx)
+func GetDatabricksProfile(ctx context.Context) string {
+	sess, err := session.GetSession(ctx)
+	if err != nil {
+		return ""
+	}
+	profile, ok := sess.Get(DatabricksProfileKey)
+	if !ok {
+		return ""
+	}
+	return profile.(string)
+}
+
+// GetAvailableProfiles returns all available profiles from ~/.databrickscfg.
+func GetAvailableProfiles(ctx context.Context) profile.Profiles {
+	profiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.MatchAllProfiles)
+	if err != nil {
+		// If we can't load profiles, return empty list (config file might not exist)
+		return profile.Profiles{}
+	}
+	return profiles
+}
+
+func MustGetApiClient(ctx context.Context) *httpclient.ApiClient {
+	client, err := GetApiClient(ctx)
+	if err != nil {
+		panic(err)
+	}
+	return client
+}
+
+func GetApiClient(ctx context.Context) (*httpclient.ApiClient, error) {
+	w, err := GetDatabricksClient(ctx)
+	if err != nil {
+		return nil, err
+	}
 	clientCfg, err := config.HTTPClientConfigFromConfig(w.Config)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create HTTP client config: %w", err)
@@ -64,28 +99,36 @@ func GetDatabricksClient(ctx context.Context) (*databricks.WorkspaceClient, erro
 	}
 	w, ok := sess.Get(DatabricksClientKey)
 	if !ok {
-		return nil, errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", nil))
+		return nil, newAuthError(ctx)
 	}
 	return w.(*databricks.WorkspaceClient), nil
 }
 
 func checkAuth(ctx context.Context) (*databricks.WorkspaceClient, error) {
 	w, err := databricks.NewWorkspaceClient()
 	if err != nil {
-		return nil, wrapAuthError(err)
+		return nil, WrapAuthError(ctx, err)
 	}
 
 	_, err = w.CurrentUser.Me(ctx)
 	if err != nil {
-		return nil, wrapAuthError(err)
+		return nil, WrapAuthError(ctx, err)
 	}
 
 	return w, nil
 }
 
-func wrapAuthError(err error) error {
+func WrapAuthError(ctx context.Context, err error) error {
 	if errors.Is(err, config.ErrCannotConfigureDefault) {
-		return errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", nil))
+		return newAuthError(ctx)
 	}
 	return err
 }
+
+func newAuthError(ctx context.Context) error {
+	// Prepare template data
+	data := map[string]any{
+		"Profiles": GetAvailableProfiles(ctx),
+	}
+	return errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", data))
+}

experimental/apps-mcp/lib/middlewares/warehouse.go

@@ -86,7 +86,10 @@ func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) {
 	// first resolve DATABRICKS_WAREHOUSE_ID env variable
 	warehouseID := env.Get(ctx, "DATABRICKS_WAREHOUSE_ID")
 	if warehouseID != "" {
-		w := MustGetDatabricksClient(ctx)
+		w, err := GetDatabricksClient(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("get databricks client: %w", err)
+		}
 		warehouse, err := w.Warehouses.Get(ctx, sql.GetWarehouseRequest{
 			Id: warehouseID,
 		})
@@ -100,7 +103,7 @@ func getDefaultWarehouse(ctx context.Context) (*sql.EndpointInfo, error) {
 		}, nil
 	}
 
-	apiClient, err := MustGetApiClient(ctx)
+	apiClient, err := GetApiClient(ctx)
 	if err != nil {
 		return nil, err
 	}

experimental/apps-mcp/lib/prompts/auth_error.tmpl

@@ -6,9 +6,14 @@
 Not authenticated to Databricks
 
 I need to know either the Databricks workspace URL or the Databricks profile name.
-You can list the available profiles by running `databricks auth profiles`.
 
-ASK the user which of the configured profiles or databricks workspace URL they want to use.
+The available profiles are:
+
+{{- range .Profiles }}
+- {{ .Name }} ({{ .Host }})
+{{- end }}
+
+IMPORTANT: YOU MUSTASK the user which of the configured profiles or databricks workspace URL they want to use.
 Only then call the `databricks_configure_auth` tool to configure the authentication.
 
 Do not run anything else before authenticating successfully.

experimental/apps-mcp/lib/providers/clitools/configure_auth.go

@@ -9,7 +9,6 @@ import (
 	"github.com/databricks/cli/experimental/apps-mcp/lib/prompts"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/session"
 	"github.com/databricks/databricks-sdk-go"
-	"github.com/databricks/databricks-sdk-go/config"
 )
 
 // ConfigureAuth creates and validates a Databricks workspace client with optional host and profile.
@@ -20,9 +19,8 @@ func ConfigureAuth(ctx context.Context, sess *session.Session, host, profile *st
 		return nil, nil
 	}
 
-	var cfg *databricks.Config
+	cfg := &databricks.Config{}
 	if host != nil || profile != nil {
-		cfg = &databricks.Config{}
 		if host != nil {
 			cfg.Host = *host
 		}
@@ -32,12 +30,7 @@ func ConfigureAuth(ctx context.Context, sess *session.Session, host, profile *st
 	}
 
 	var client *databricks.WorkspaceClient
-	var err error
-	if cfg != nil {
-		client, err = databricks.NewWorkspaceClient(cfg)
-	} else {
-		client, err = databricks.NewWorkspaceClient()
-	}
+	client, err := databricks.NewWorkspaceClient(cfg)
 	if err != nil {
 		return nil, err
 	}
@@ -49,19 +42,14 @@ func ConfigureAuth(ctx context.Context, sess *session.Session, host, profile *st
 				"WorkspaceURL": *host,
 			}))
 		}
-		return nil, wrapAuthError(err)
+		return nil, middlewares.WrapAuthError(ctx, err)
 	}
 
-	// Store client in session data
+	// Store client and profile in session data
 	sess.Set(middlewares.DatabricksClientKey, client)
+	if profile != nil {
+		sess.Set(middlewares.DatabricksProfileKey, *profile)
+	}
 
 	return client, nil
 }
-
-// wrapAuthError wraps configuration errors with helpful messages
-func wrapAuthError(err error) error {
-	if errors.Is(err, config.ErrCannotConfigureDefault) {
-		return errors.New(prompts.MustExecuteTemplate("auth_error.tmpl", nil))
-	}
-	return err
-}

experimental/apps-mcp/lib/providers/clitools/explore.go

@@ -8,7 +8,6 @@ import (
 	"github.com/databricks/cli/experimental/apps-mcp/lib/prompts"
 	"github.com/databricks/cli/experimental/apps-mcp/lib/session"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
-	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 )
@@ -21,32 +20,12 @@ func Explore(ctx context.Context) (string, error) {
 		warehouse = nil
 	}
 
-	currentProfile := getCurrentProfile(ctx)
-	profiles := getAvailableProfiles(ctx)
+	currentProfile := middlewares.GetDatabricksProfile(ctx)
+	profiles := middlewares.GetAvailableProfiles(ctx)
 
 	return generateExploreGuidance(ctx, warehouse, currentProfile, profiles), nil
 }
 
-// getCurrentProfile returns the currently active profile name.
-func getCurrentProfile(ctx context.Context) string {
-	// Check DATABRICKS_CONFIG_PROFILE env var
-	profileName := env.Get(ctx, "DATABRICKS_CONFIG_PROFILE")
-	if profileName == "" {
-		return "DEFAULT"
-	}
-	return profileName
-}
-
-// getAvailableProfiles returns all available profiles from ~/.databrickscfg.
-func getAvailableProfiles(ctx context.Context) profile.Profiles {
-	profiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.MatchAllProfiles)
-	if err != nil {
-		// If we can't load profiles, return empty list (config file might not exist)
-		return profile.Profiles{}
-	}
-	return profiles
-}
-
 // generateExploreGuidance creates comprehensive guidance for data exploration.
 func generateExploreGuidance(ctx context.Context, warehouse *sql.EndpointInfo, currentProfile string, profiles profile.Profiles) string {
 	// Build workspace/profile information

experimental/apps-mcp/lib/providers/clitools/invoke_databricks_cli.go

@@ -17,15 +17,22 @@ func InvokeDatabricksCLI(ctx context.Context, command []string, workingDirectory
 		return "", errors.New("command is required")
 	}
 
-	workspaceClient := middlewares.MustGetDatabricksClient(ctx)
+	workspaceClient, err := middlewares.GetDatabricksClient(ctx)
+	if err != nil {
+		return "", fmt.Errorf("get databricks client: %w", err)
+	}
 	host := workspaceClient.Config.Host
+	profile := middlewares.GetDatabricksProfile(ctx)
 
 	// GetCLIPath returns the path to the current CLI executable
 	cliPath := common.GetCLIPath()
 	cmd := exec.CommandContext(ctx, cliPath, command...)
 	cmd.Dir = workingDirectory
 	env := os.Environ()
 	env = append(env, "DATABRICKS_HOST="+host)
+	if profile != "" {
+		env = append(env, "DATABRICKS_CONFIG_PROFILE="+profile)
+	}
 	cmd.Env = env
 
 	output, err := cmd.CombinedOutput()

experimental/apps-mcp/lib/providers/clitools/provider.go

@@ -93,7 +93,7 @@ func (p *Provider) RegisterTools(server *mcpsdk.Server) error {
 		},
 		func(ctx context.Context, req *mcpsdk.CallToolRequest, args struct{}) (*mcpsdk.CallToolResult, any, error) {
 			log.Debug(ctx, "explore called")
-			result, err := Explore(session.WithSession(ctx, p.session))
+			result, err := Explore(ctx)
 			if err != nil {
 				return nil, nil, err
 			}