diff --git a/cmd/apps/deploy_bundle.go b/cmd/apps/deploy_bundle.go index 38e312e388..86b8c8bb1e 100644 --- a/cmd/apps/deploy_bundle.go +++ b/cmd/apps/deploy_bundle.go @@ -12,6 +12,7 @@ import ( "github.com/databricks/cli/bundle/run" "github.com/databricks/cli/cmd/bundle/utils" "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/apps/prompt" "github.com/databricks/cli/libs/apps/validation" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/log" @@ -153,11 +154,11 @@ func runBundleDeploy(cmd *cobra.Command, force, skipValidation, skipTests bool) log.Infof(ctx, "Running app: %s", appKey) if err := runBundleApp(ctx, b, appKey); err != nil { - cmdio.LogString(ctx, "✔ Deployment succeeded, but failed to start app") + prompt.PrintDone(ctx, "Deployment succeeded, but failed to start app") return fmt.Errorf("failed to run app: %w. Run `databricks apps logs` to view logs", err) } - cmdio.LogString(ctx, "✔ Deployment complete!") + prompt.PrintDone(ctx, "Deployment complete!") return nil } diff --git a/cmd/apps/import.go b/cmd/apps/import.go index b508335231..d32112f5c5 100644 --- a/cmd/apps/import.go +++ b/cmd/apps/import.go @@ -20,6 +20,7 @@ import ( "github.com/databricks/cli/bundle/run" bundleutils "github.com/databricks/cli/cmd/bundle/utils" "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/apps/prompt" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/dyn" @@ -202,9 +203,10 @@ Examples: } if !quiet { - cmdio.LogString(ctx, fmt.Sprintf("\n✓ App '%s' has been successfully imported to %s", name, outputDir)) + cmdio.LogString(ctx, "") + prompt.PrintDone(ctx, fmt.Sprintf("App '%s' imported to %s", name, outputDir)) if cleanup && oldSourceCodePath != "" { - cmdio.LogString(ctx, "✓ Previous app folder has been cleaned up") + prompt.PrintDone(ctx, "Previous app folder cleaned up") } cmdio.LogString(ctx, "\nYou can now deploy changes with: databricks bundle deploy") } diff --git a/cmd/apps/init.go b/cmd/apps/init.go index e04b7fcc8f..089dd48011 100644 --- a/cmd/apps/init.go +++ b/cmd/apps/init.go @@ -339,6 +339,14 @@ func promptForPluginsAndDeps(ctx context.Context, m *manifest.Manifest, preSelec } theme := prompt.AppkitTheme() + // Eagerly start fetching resources for ALL plugins in the background. + // This runs while the user is selecting plugins, so by the time resource + // pickers appear the data is likely already cached. + allPluginNames := m.GetPluginNames() + allPossibleResources := m.CollectResources(allPluginNames) + allPossibleResources = append(allPossibleResources, m.CollectOptionalResources(allPluginNames)...) + ctx = prompt.PrefetchResources(ctx, allPossibleResources) + // Step 1: Plugin selection (skip if plugins already provided via flag) selectablePlugins := m.GetSelectablePlugins() if len(config.Features) == 0 && len(selectablePlugins) > 0 { @@ -371,8 +379,11 @@ func promptForPluginsAndDeps(ctx context.Context, m *manifest.Manifest, preSelec // Always include mandatory plugins. config.Features = appendUnique(config.Features, m.GetMandatoryPluginNames()...) - // Step 2: Prompt for required plugin resource dependencies + // Collect resources for the user's actual selection. resources := m.CollectResources(config.Features) + optionalResources := m.CollectOptionalResources(config.Features) + + // Step 2: Prompt for required plugin resource dependencies for _, r := range resources { values, err := promptForResource(ctx, r, theme, true) if err != nil { @@ -384,7 +395,6 @@ func promptForPluginsAndDeps(ctx context.Context, m *manifest.Manifest, preSelec } // Step 3: Prompt for optional plugin resource dependencies - optionalResources := m.CollectOptionalResources(config.Features) for _, r := range optionalResources { values, err := promptForResource(ctx, r, theme, false) if err != nil { @@ -500,43 +510,193 @@ func cloneRepo(ctx context.Context, repoURL, branch string) (string, error) { return tempDir, nil } -// resolveTemplate resolves a template path, handling both local paths and GitHub URLs. -// branch is used for cloning (can contain "/" for feature branches). -// subdir is an optional subdirectory within the repo to use (for default appkit template). -// Returns the local path to use, a cleanup function (for temp dirs), and any error. -func resolveTemplate(ctx context.Context, templatePath, branch, subdir string) (localPath string, cleanup func(), err error) { - // Case 1: Local path - return as-is +// resolveTemplate resolves a template synchronously with a spinner. +// Used by commands that don't benefit from background cloning (e.g., manifest). +func resolveTemplate(ctx context.Context, templatePath, branch, subdir string) (string, func(), error) { + ch := resolveTemplateAsync(ctx, templatePath, branch, subdir) + return awaitTemplate(ctx, ch) +} + +// templateResult holds the outcome of a background template resolution. +type templateResult struct { + path string + cleanup func() + err error +} + +// resolveTemplateAsync starts resolving the template in a background goroutine. +// For local paths this completes immediately; for GitHub URLs it clones the repo. +// The caller reads the result from the returned channel, optionally showing a +// spinner if the clone hasn't finished by the time it's needed. +func resolveTemplateAsync(ctx context.Context, templatePath, branch, subdir string) <-chan templateResult { + ch := make(chan templateResult, 1) + + // Local path — instant. if !strings.HasPrefix(templatePath, "https://") { - return templatePath, nil, nil + ch <- templateResult{path: templatePath} + return ch } - // Case 2: GitHub URL - parse and clone repoURL, urlSubdir, urlBranch := git.ParseGitHubURL(templatePath) if branch == "" { - branch = urlBranch // Use branch from URL if not overridden by flag + branch = urlBranch } if subdir == "" { - subdir = urlSubdir // Use subdir from URL if not overridden + subdir = urlSubdir } - // Clone to temp dir with spinner - var tempDir string - err = prompt.RunWithSpinnerCtx(ctx, "Cloning template...", func() error { - var cloneErr error - tempDir, cloneErr = cloneRepo(ctx, repoURL, branch) - return cloneErr - }) + go func() { + tempDir, err := cloneRepo(ctx, repoURL, branch) + if err != nil { + ch <- templateResult{err: err} + return + } + cleanup := func() { os.RemoveAll(tempDir) } + localPath := tempDir + if subdir != "" { + localPath = filepath.Join(tempDir, subdir) + } + ch <- templateResult{path: localPath, cleanup: cleanup} + }() + + return ch +} + +// awaitTemplate waits for the background clone to finish. +// If the result is already available it returns immediately with a +// checkmark; otherwise it shows a spinner while waiting. +func awaitTemplate(ctx context.Context, ch <-chan templateResult) (string, func(), error) { + select { + case res := <-ch: + // Clone finished while the user was typing — print completion. + if res.err == nil && res.cleanup != nil { + prompt.PrintDone(ctx, "Template cloned") + } + return res.path, res.cleanup, res.err + default: + // Still cloning — show a spinner for the remaining wait. + var res templateResult + err := prompt.RunWithSpinnerCtx(ctx, "Cloning template...", func() error { + res = <-ch + return res.err + }) + return res.path, res.cleanup, err + } +} + +// findProjectSrcDir locates the actual source directory inside a template. +// Templates may nest their content inside a {{.project_name}} directory. +func findProjectSrcDir(templateDir string) string { + entries, err := os.ReadDir(templateDir) if err != nil { - return "", nil, err + return templateDir + } + for _, e := range entries { + if e.IsDir() && strings.Contains(e.Name(), "{{.project_name}}") { + return filepath.Join(templateDir, e.Name()) + } } + return templateDir +} - cleanup = func() { os.RemoveAll(tempDir) } +// startBackgroundNpmInstall copies the package files from the template into +// destDir and launches `npm ci` in the background. The caller should await +// the returned channel BEFORE writing other files to destDir to prevent +// concurrent writes. Returns nil if the template is not a Node.js project +// or npm is not available. +// +// IMPORTANT: All reads from srcProjectDir happen synchronously before the +// goroutine launches. The template directory may be cleaned up after this +// function returns, so file reads must not be deferred to the goroutine. +func startBackgroundNpmInstall(ctx context.Context, srcProjectDir, destDir, projectName string) <-chan error { + lockFile := filepath.Join(srcProjectDir, "package-lock.json") + if _, err := os.Stat(lockFile); err != nil { + return nil + } + + if _, err := exec.LookPath("npm"); err != nil { + return nil + } - // Return path to subdirectory if specified - if subdir != "" { - return filepath.Join(tempDir, subdir), cleanup, nil + if err := os.MkdirAll(destDir, 0o755); err != nil { + log.Warnf(ctx, "Failed to create %s: %v, skipping background npm install", destDir, err) + return nil + } + + // Copy package.json (apply template substitution so the file is valid JSON) + // and package-lock.json (no template vars — copy raw). + var pkgWritten bool + for _, name := range []string{"package.json", "package.json.tmpl"} { + src := filepath.Join(srcProjectDir, name) + content, err := os.ReadFile(src) + if err != nil { + continue + } + minVars := templateData(templateVars{ + ProjectName: projectName, + AppDescription: prompt.DefaultAppDescription, + Plugins: make(map[string]*pluginVar), + }) + tmpl, err := template.New(name).Option("missingkey=zero").Parse(string(content)) + if err != nil { + pkgWritten = os.WriteFile(filepath.Join(destDir, "package.json"), content, 0o644) == nil + break + } + var buf bytes.Buffer + if err := tmpl.Execute(&buf, minVars); err != nil { + pkgWritten = os.WriteFile(filepath.Join(destDir, "package.json"), content, 0o644) == nil + break + } + pkgWritten = os.WriteFile(filepath.Join(destDir, "package.json"), buf.Bytes(), 0o644) == nil + break + } + + if !pkgWritten { + log.Warnf(ctx, "Failed to write package.json to %s, skipping background npm install", destDir) + return nil + } + + // Copy package-lock.json raw (never has template vars). + lockData, err := os.ReadFile(lockFile) + if err != nil { + log.Warnf(ctx, "Failed to read package-lock.json: %v, skipping background npm install", err) + return nil + } + if err := os.WriteFile(filepath.Join(destDir, "package-lock.json"), lockData, 0o644); err != nil { + log.Warnf(ctx, "Failed to write package-lock.json: %v, skipping background npm install", err) + return nil + } + + ch := make(chan error, 1) + go func() { + cmd := exec.CommandContext(ctx, "npm", "ci", "--no-audit", "--no-fund", "--prefer-offline") + cmd.Dir = destDir + cmd.Stdout = nil + cmd.Stderr = nil + ch <- cmd.Run() + }() + + log.Debugf(ctx, "Started background npm install in %s", destDir) + return ch +} + +// awaitBackgroundNpmInstall waits for the background npm install to complete. +// Shows an instant checkmark if already done, or a spinner for the remainder. +func awaitBackgroundNpmInstall(ctx context.Context, ch <-chan error) error { + select { + case err := <-ch: + if err == nil { + prompt.PrintDone(ctx, "Dependencies installed") + } + return err + default: + var installErr error + err := prompt.RunWithSpinnerCtx(ctx, "Installing dependencies...", func() error { + installErr = <-ch + return installErr + }) + return err } - return tempDir, cleanup, nil } func runCreate(ctx context.Context, opts createOptions) error { @@ -574,8 +734,25 @@ func runCreate(ctx context.Context, opts createOptions) error { templateSrc = appkitRepoURL } - // Step 1: Get project name first (needed before we can check destination) - // Determine output directory for validation + // Start cloning in the background so it runs while the user types the name. + branchForClone := opts.branch + subdirForClone := "" + if usingDefaultTemplate { + branchForClone = gitRef + subdirForClone = appkitTemplateDir + } + templateCh := resolveTemplateAsync(ctx, templateSrc, branchForClone, subdirForClone) + defer func() { + select { + case res := <-templateCh: + if res.cleanup != nil { + res.cleanup() + } + default: + } + }() + + // Step 1: Get project name (clone runs in parallel for remote templates) destDir := opts.name if opts.outputDir != "" { destDir = filepath.Join(opts.outputDir, opts.name) @@ -585,19 +762,16 @@ func runCreate(ctx context.Context, opts createOptions) error { if !isInteractive { return errors.New("--name is required in non-interactive mode") } - // Prompt includes validation for name format AND directory existence name, err := prompt.PromptForProjectName(ctx, opts.outputDir) if err != nil { return err } opts.name = name - // Update destDir with the actual name destDir = opts.name if opts.outputDir != "" { destDir = filepath.Join(opts.outputDir, opts.name) } } else { - // Non-interactive mode: validate name and directory existence if err := prompt.ValidateProjectName(opts.name); err != nil { return err } @@ -606,16 +780,8 @@ func runCreate(ctx context.Context, opts createOptions) error { } } - // Step 2: Resolve template (handles GitHub URLs by cloning) - // For custom templates, --branch can override the URL's branch - // For default appkit template, pass gitRef directly (supports branches with "/" in name) - branchForClone := opts.branch - subdirForClone := "" - if usingDefaultTemplate { - branchForClone = gitRef - subdirForClone = appkitTemplateDir - } - resolvedPath, cleanup, err := resolveTemplate(ctx, templateSrc, branchForClone, subdirForClone) + // Step 2: Wait for template (may already be done if the user took time typing the name) + resolvedPath, cleanup, err := awaitTemplate(ctx, templateCh) if err != nil { return err } @@ -633,6 +799,11 @@ func runCreate(ctx context.Context, opts createOptions) error { } } + // Start npm install in the background so it runs while the user answers prompts. + // This is a Node.js-only optimisation — non-Node templates skip this. + srcProjectDir := findProjectSrcDir(templateDir) + npmInstallCh := startBackgroundNpmInstall(ctx, srcProjectDir, destDir, opts.name) + // Step 3: Load manifest from template (optional — templates without it skip plugin/resource logic) var m *manifest.Manifest if manifest.HasManifest(templateDir) { @@ -770,12 +941,12 @@ func runCreate(ctx context.Context, opts createOptions) error { } } - // Track whether we started creating the project for cleanup on failure + // Track whether we started creating the project for cleanup on failure. + // The background npm install may have created destDir early. var projectCreated bool var runErr error defer func() { - if runErr != nil && projectCreated { - // Clean up partially created project on failure + if runErr != nil && (projectCreated || npmInstallCh != nil) { os.RemoveAll(destDir) } }() @@ -841,6 +1012,17 @@ func runCreate(ctx context.Context, opts createOptions) error { Plugins: plugins, } + // Await background npm install BEFORE copying the template so there are + // no concurrent writes to destDir. npm ci ran with the raw lock file; the + // dependency tree is determined entirely by package-lock.json which has no + // template variables, so the installed node_modules is valid. + if npmInstallCh != nil { + if err := awaitBackgroundNpmInstall(ctx, npmInstallCh); err != nil { + log.Warnf(ctx, "Background npm install failed: %v, will retry during project initialization", err) + os.RemoveAll(filepath.Join(destDir, "node_modules")) + } + } + // Copy template with variable substitution var fileCount int runErr = prompt.RunWithSpinnerCtx(ctx, "Creating project...", func() error { @@ -859,7 +1041,9 @@ func runCreate(ctx context.Context, opts createOptions) error { absOutputDir = destDir } - // Initialize project based on type (Node.js, Python, etc.) + // Initialize project based on type (Node.js, Python, etc.). + // For Node.js, if the background install succeeded node_modules exists + // and the initializer skips the redundant install step. var nextStepsCmd string projectInitializer := initializer.GetProjectInitializer(absOutputDir) if projectInitializer != nil { @@ -925,10 +1109,11 @@ func runCreate(ctx context.Context, opts createOptions) error { if shouldDeploy { cmdio.LogString(ctx, "") - cmdio.LogString(ctx, "Deploying app...") if err := runPostCreateDeploy(ctx, profile); err != nil { cmdio.LogString(ctx, fmt.Sprintf("⚠ Deploy failed: %v", err)) cmdio.LogString(ctx, " You can deploy manually with: databricks apps deploy") + } else { + prompt.PrintDone(ctx, "Deploy complete") } } diff --git a/libs/apps/generator/generator_test.go b/libs/apps/generator/generator_test.go index fd4f85cf45..a3a7946fbc 100644 --- a/libs/apps/generator/generator_test.go +++ b/libs/apps/generator/generator_test.go @@ -1040,3 +1040,45 @@ func TestBundleIgnoreFieldSkippedInVariablesAndTargets(t *testing.T) { assert.Contains(t, example, "DB_INSTANCE=your_database_instance_name") assert.Contains(t, example, "DB_NAME=your_database_database_name") } + +func TestVolumeManifestPathFieldMapsToSpecId(t *testing.T) { + plugins := []manifest.Plugin{ + { + Name: "files", + Resources: manifest.Resources{ + Required: []manifest.Resource{ + { + Type: "volume", + Alias: "Files", + ResourceKey: "files", + Description: "Permission to write to volumes", + Permission: "WRITE_VOLUME", + Fields: map[string]manifest.ResourceField{ + "path": {Env: "DATABRICKS_VOLUME_FILES", Description: "Volume path"}, + }, + }, + }, + }, + }, + } + + cfg := generator.Config{ + ResourceValues: map[string]string{ + "files.path": "/Volumes/catalog/schema/vol", + "files.id": "catalog.schema.vol", + }, + } + + vars := generator.GenerateBundleVariables(plugins, cfg) + assert.Contains(t, vars, "files_path:") + assert.Contains(t, vars, "files_id:") + + target := generator.GenerateTargetVariables(plugins, cfg) + assert.Contains(t, target, "files_path: /Volumes/catalog/schema/vol") + assert.Contains(t, target, "files_id: catalog.schema.vol") + + res := generator.GenerateBundleResources(plugins, cfg) + assert.Contains(t, res, "securable_full_name: ${var.files_id}") + assert.Contains(t, res, "securable_type: VOLUME") + assert.Contains(t, res, "permission: WRITE_VOLUME") +} diff --git a/libs/apps/initializer/nodejs.go b/libs/apps/initializer/nodejs.go index 1e96f43f72..958f533558 100644 --- a/libs/apps/initializer/nodejs.go +++ b/libs/apps/initializer/nodejs.go @@ -20,12 +20,14 @@ type InitializerNodeJs struct { func (i *InitializerNodeJs) Initialize(ctx context.Context, workDir string) *InitResult { i.workDir = workDir - // Step 1: Run npm install - if err := i.runNpmInstall(ctx, workDir); err != nil { - return &InitResult{ - Success: false, - Message: "Failed to install dependencies", - Error: err, + // Step 1: Run npm install (skip if node_modules already exists from a background install) + if !fileExists(filepath.Join(workDir, "node_modules")) { + if err := i.runNpmInstall(ctx, workDir); err != nil { + return &InitResult{ + Success: false, + Message: "Failed to install dependencies", + Error: err, + } } } diff --git a/libs/apps/manifest/manifest.go b/libs/apps/manifest/manifest.go index 21ea16c8e4..84f3018164 100644 --- a/libs/apps/manifest/manifest.go +++ b/libs/apps/manifest/manifest.go @@ -30,6 +30,10 @@ type Resource struct { Description string `json:"description"` // e.g., "SQL Warehouse for executing analytics queries" Permission string `json:"permission"` // e.g., "CAN_USE" Fields map[string]ResourceField `json:"fields"` // field definitions with env var mappings + + // PluginDisplayName is set during resource collection to identify which + // plugin requires this resource. Not part of the JSON manifest. + PluginDisplayName string `json:"-"` } // Key returns the resource key for machine use (config keys, variable naming). @@ -185,6 +189,7 @@ func (m *Manifest) ValidatePluginNames(names []string) error { } // CollectResources returns all required resources for the given plugin names. +// Each returned resource is annotated with PluginDisplayName for UI context. func (m *Manifest) CollectResources(pluginNames []string) []Resource { seen := make(map[string]bool) var resources []Resource @@ -202,6 +207,7 @@ func (m *Manifest) CollectResources(pluginNames []string) []Resource { key := r.Type + ":" + r.Key() if !seen[key] { seen[key] = true + r.PluginDisplayName = plugin.DisplayName resources = append(resources, r) } } @@ -211,6 +217,7 @@ func (m *Manifest) CollectResources(pluginNames []string) []Resource { } // CollectOptionalResources returns all optional resources for the given plugin names. +// Each returned resource is annotated with PluginDisplayName for UI context. func (m *Manifest) CollectOptionalResources(pluginNames []string) []Resource { seen := make(map[string]bool) var resources []Resource @@ -228,6 +235,7 @@ func (m *Manifest) CollectOptionalResources(pluginNames []string) []Resource { key := r.Type + ":" + r.Key() if !seen[key] { seen[key] = true + r.PluginDisplayName = plugin.DisplayName resources = append(resources, r) } } diff --git a/libs/apps/prompt/cache.go b/libs/apps/prompt/cache.go new file mode 100644 index 0000000000..a81e5b88c6 --- /dev/null +++ b/libs/apps/prompt/cache.go @@ -0,0 +1,47 @@ +package prompt + +import ( + "context" + "sync" +) + +type resourceCacheKey struct{} + +// ResourceCache stores pre-fetched PagedFetcher instances so prompt functions +// can skip the initial API fetch when data has already been loaded in parallel. +type ResourceCache struct { + mu sync.RWMutex + fetchers map[string]*PagedFetcher +} + +// NewResourceCache creates an empty cache. +func NewResourceCache() *ResourceCache { + return &ResourceCache{fetchers: make(map[string]*PagedFetcher)} +} + +// SetFetcher stores a pre-created PagedFetcher for a resource type. +func (c *ResourceCache) SetFetcher(resourceType string, f *PagedFetcher) { + c.mu.Lock() + defer c.mu.Unlock() + c.fetchers[resourceType] = f +} + +// GetFetcher returns the cached PagedFetcher for a resource type, or nil. +func (c *ResourceCache) GetFetcher(resourceType string) *PagedFetcher { + c.mu.RLock() + defer c.mu.RUnlock() + return c.fetchers[resourceType] +} + +// ContextWithCache returns a child context carrying the given resource cache. +func ContextWithCache(ctx context.Context, cache *ResourceCache) context.Context { + return context.WithValue(ctx, resourceCacheKey{}, cache) +} + +// CacheFromContext retrieves the resource cache from the context, or nil. +func CacheFromContext(ctx context.Context) *ResourceCache { + if cache, ok := ctx.Value(resourceCacheKey{}).(*ResourceCache); ok { + return cache + } + return nil +} diff --git a/libs/apps/prompt/cache_test.go b/libs/apps/prompt/cache_test.go new file mode 100644 index 0000000000..c882dd3ee2 --- /dev/null +++ b/libs/apps/prompt/cache_test.go @@ -0,0 +1,69 @@ +package prompt + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestResourceCacheSetAndGet(t *testing.T) { + cache := NewResourceCache() + f := &PagedFetcher{Items: []ListItem{{ID: "1", Label: "one"}}} + + cache.SetFetcher("sql_warehouse", f) + got := cache.GetFetcher("sql_warehouse") + require.NotNil(t, got) + assert.Equal(t, f, got) +} + +func TestResourceCacheGetMissReturnsNil(t *testing.T) { + cache := NewResourceCache() + assert.Nil(t, cache.GetFetcher("nonexistent")) +} + +func TestResourceCacheOverwrite(t *testing.T) { + cache := NewResourceCache() + f1 := &PagedFetcher{Items: []ListItem{{ID: "1"}}} + f2 := &PagedFetcher{Items: []ListItem{{ID: "2"}}} + + cache.SetFetcher("key", f1) + cache.SetFetcher("key", f2) + + got := cache.GetFetcher("key") + assert.Equal(t, f2, got) +} + +func TestResourceCacheConcurrentAccess(t *testing.T) { + cache := NewResourceCache() + var wg sync.WaitGroup + + for i := range 100 { + wg.Go(func() { + key := "key" + f := &PagedFetcher{Items: []ListItem{{ID: string(rune('A' + i%26))}}} + cache.SetFetcher(key, f) + _ = cache.GetFetcher(key) + }) + } + + wg.Wait() + assert.NotNil(t, cache.GetFetcher("key")) +} + +func TestContextWithCacheRoundTrip(t *testing.T) { + ctx := t.Context() + cache := NewResourceCache() + + assert.Nil(t, CacheFromContext(ctx), "no cache in bare context") + + ctx = ContextWithCache(ctx, cache) + got := CacheFromContext(ctx) + require.NotNil(t, got) + assert.Equal(t, cache, got) +} + +func TestCacheFromContextNilOnMissingValue(t *testing.T) { + assert.Nil(t, CacheFromContext(t.Context())) +} diff --git a/libs/apps/prompt/listers.go b/libs/apps/prompt/listers.go index a78c28962a..5ccf818e01 100644 --- a/libs/apps/prompt/listers.go +++ b/libs/apps/prompt/listers.go @@ -20,6 +20,7 @@ import ( "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/databricks/databricks-sdk-go/service/ml" "github.com/databricks/databricks-sdk-go/service/postgres" + "github.com/databricks/databricks-sdk-go/service/serving" "github.com/databricks/databricks-sdk-go/service/sql" "github.com/databricks/databricks-sdk-go/service/vectorsearch" "github.com/databricks/databricks-sdk-go/service/workspace" @@ -87,29 +88,6 @@ func ListSecretKeys(ctx context.Context, scope string) ([]ListItem, error) { return out, nil } -// ListJobs returns jobs as selectable items. -func ListJobs(ctx context.Context) ([]ListItem, error) { - w, err := workspaceClient(ctx) - if err != nil { - return nil, err - } - iter := w.Jobs.List(ctx, jobs.ListJobsRequest{}) - jobList, err := listing.ToSlice(ctx, iter) - if err != nil { - return nil, err - } - out := make([]ListItem, 0, len(jobList)) - for _, j := range jobList { - label := j.Settings.Name - id := strconv.FormatInt(j.JobId, 10) - if label == "" { - label = id - } - out = append(out, ListItem{ID: id, Label: label}) - } - return out, nil -} - // ListSQLWarehousesItems returns SQL warehouses as ListItems (reuses same API as ListSQLWarehouses). func ListSQLWarehousesItems(ctx context.Context) ([]ListItem, error) { w, err := workspaceClient(ctx) @@ -121,7 +99,7 @@ func ListSQLWarehousesItems(ctx context.Context) ([]ListItem, error) { if err != nil { return nil, err } - out := make([]ListItem, 0, len(whs)) + out := make([]ListItem, 0, min(len(whs), maxListResults)) for _, wh := range whs { label := wh.Name if wh.State != "" { @@ -129,46 +107,6 @@ func ListSQLWarehousesItems(ctx context.Context) ([]ListItem, error) { } out = append(out, ListItem{ID: wh.Id, Label: label}) } - return out, nil -} - -// ListServingEndpoints returns serving endpoints as selectable items. -func ListServingEndpoints(ctx context.Context) ([]ListItem, error) { - w, err := workspaceClient(ctx) - if err != nil { - return nil, err - } - iter := w.ServingEndpoints.List(ctx) - endpoints, err := listing.ToSlice(ctx, iter) - if err != nil { - return nil, err - } - out := make([]ListItem, 0, len(endpoints)) - for _, e := range endpoints { - name := e.Name - if name == "" { - name = e.Id - } - out = append(out, ListItem{ID: e.Id, Label: name}) - } - return out, nil -} - -// ListCatalogs returns UC catalogs as selectable items. -func ListCatalogs(ctx context.Context) ([]ListItem, error) { - w, err := workspaceClient(ctx) - if err != nil { - return nil, err - } - catIter := w.Catalogs.List(ctx, catalog.ListCatalogsRequest{}) - cats, err := listing.ToSlice(ctx, catIter) - if err != nil { - return nil, err - } - out := make([]ListItem, 0, min(len(cats), maxListResults)) - for _, c := range cats { - out = append(out, ListItem{ID: c.Name, Label: c.Name}) - } return capResults(out), nil } @@ -206,43 +144,12 @@ func ListVolumesInSchema(ctx context.Context, catalogName, schemaName string) ([ } out := make([]ListItem, 0, min(len(vols), maxListResults)) for _, v := range vols { - fullName := fmt.Sprintf("%s.%s.%s", catalogName, schemaName, v.Name) - out = append(out, ListItem{ID: fullName, Label: v.Name}) + volumePath := fmt.Sprintf("/Volumes/%s/%s/%s", catalogName, schemaName, v.Name) + out = append(out, ListItem{ID: volumePath, Label: v.Name}) } return capResults(out), nil } -// ListVectorSearchIndexes returns vector search indexes as selectable items (id = endpoint/index name). -func ListVectorSearchIndexes(ctx context.Context) ([]ListItem, error) { - w, err := workspaceClient(ctx) - if err != nil { - return nil, err - } - var out []ListItem - epIter := w.VectorSearchEndpoints.ListEndpoints(ctx, vectorsearch.ListEndpointsRequest{}) - endpoints, err := listing.ToSlice(ctx, epIter) - if err != nil { - return nil, err - } - for _, ep := range endpoints { - indexIter := w.VectorSearchIndexes.ListIndexes(ctx, vectorsearch.ListIndexesRequest{EndpointName: ep.Name}) - indexes, err := listing.ToSlice(ctx, indexIter) - if err != nil { - log.Warnf(ctx, "Failed to list indexes for endpoint %q: %v", ep.Name, err) - continue - } - for _, idx := range indexes { - label := idx.Name - if label == "" { - label = ep.Name + "/ (unnamed)" - } - id := ep.Name + "/" + idx.Name - out = append(out, ListItem{ID: id, Label: fmt.Sprintf("%s / %s", ep.Name, label)}) - } - } - return out, nil -} - // ListFunctionsInSchema returns UC functions within a catalog.schema as selectable items. func ListFunctionsInSchema(ctx context.Context, catalogName, schemaName string) ([]ListItem, error) { w, err := workspaceClient(ctx) @@ -268,46 +175,6 @@ func ListFunctionsInSchema(ctx context.Context, catalogName, schemaName string) return capResults(out), nil } -// ListConnections returns UC connections as selectable items. -func ListConnections(ctx context.Context) ([]ListItem, error) { - w, err := workspaceClient(ctx) - if err != nil { - return nil, err - } - iter := w.Connections.List(ctx, catalog.ListConnectionsRequest{}) - conns, err := listing.ToSlice(ctx, iter) - if err != nil { - return nil, err - } - out := make([]ListItem, 0, len(conns)) - for _, c := range conns { - name := c.Name - if name == "" { - name = c.FullName - } - out = append(out, ListItem{ID: c.FullName, Label: name}) - } - return out, nil -} - -// ListDatabaseInstances returns Lakebase database instances as selectable items. -func ListDatabaseInstances(ctx context.Context) ([]ListItem, error) { - w, err := workspaceClient(ctx) - if err != nil { - return nil, err - } - iter := w.Database.ListDatabaseInstances(ctx, database.ListDatabaseInstancesRequest{}) - instances, err := listing.ToSlice(ctx, iter) - if err != nil { - return nil, err - } - out := make([]ListItem, 0, len(instances)) - for _, inst := range instances { - out = append(out, ListItem{ID: inst.Name, Label: inst.Name}) - } - return out, nil -} - // listDatabasesResponse is the response from the /databases endpoint. type listDatabasesResponse struct { Databases []struct { @@ -356,28 +223,6 @@ func extractIDFromName(name, component string) string { return name } -// ListPostgresProjects returns Lakebase Autoscaling (V2) projects as selectable items. -func ListPostgresProjects(ctx context.Context) ([]ListItem, error) { - w, err := workspaceClient(ctx) - if err != nil { - return nil, err - } - iter := w.Postgres.ListProjects(ctx, postgres.ListProjectsRequest{}) - projects, err := listing.ToSlice(ctx, iter) - if err != nil { - return nil, err - } - out := make([]ListItem, 0, len(projects)) - for _, p := range projects { - label := p.Name - if p.Status != nil && p.Status.DisplayName != "" { - label = p.Status.DisplayName - } - out = append(out, ListItem{ID: p.Name, Label: label}) - } - return out, nil -} - // ListPostgresBranches returns branches within a Lakebase Autoscaling project as selectable items. func ListPostgresBranches(ctx context.Context, projectName string) ([]ListItem, error) { w, err := workspaceClient(ctx) @@ -441,19 +286,158 @@ func ListPostgresEndpoints(ctx context.Context, branchName string) ([]postgres.E return listing.ToSlice(ctx, iter) } -// ListGenieSpaces returns Genie spaces as selectable items. -func ListGenieSpaces(ctx context.Context) ([]ListItem, error) { +// --------------------------------------------------------------------------- +// Paged lister constructors — return a PagedFetcher that loads pageSize items +// at a time, keeping the SDK iterator alive for incremental "Load more". +// --------------------------------------------------------------------------- + +// ListSQLWarehouses lists SQL warehouses as a paged result. +func ListSQLWarehouses(ctx context.Context) (*PagedFetcher, error) { w, err := workspaceClient(ctx) if err != nil { return nil, err } - var out []ListItem - req := dashboards.GenieListSpacesRequest{} - for { + iter := w.Warehouses.List(ctx, sql.ListWarehousesRequest{PageSize: pageSize}) + mapFn := func(wh sql.EndpointInfo) ListItem { + label := wh.Name + if wh.State != "" { + label = fmt.Sprintf("%s (%s)", wh.Name, wh.State) + } + return ListItem{ID: wh.Id, Label: label} + } + items, hasMore, err := collectN(ctx, iter, pageSize, mapFn) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return collectN(ctx, iter, pageSize, mapFn) + }, + }, nil +} + +// ListJobs lists jobs as a paged result. +func ListJobs(ctx context.Context) (*PagedFetcher, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + apiLimit := min(pageSize, 100) + iter := w.Jobs.List(ctx, jobs.ListJobsRequest{Limit: apiLimit}) + mapFn := func(j jobs.BaseJob) ListItem { + label := j.Settings.Name + id := strconv.FormatInt(j.JobId, 10) + if label == "" { + label = id + } + return ListItem{ID: id, Label: label} + } + items, hasMore, err := collectN(ctx, iter, apiLimit, mapFn) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return collectN(ctx, iter, apiLimit, mapFn) + }, + }, nil +} + +// SearchJobs performs a server-side search for jobs by name (exact, case-insensitive). +func SearchJobs(ctx context.Context, name string) ([]ListItem, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + iter := w.Jobs.List(ctx, jobs.ListJobsRequest{Name: name}) + jobList, err := listing.ToSlice(ctx, iter) + if err != nil { + return nil, err + } + out := make([]ListItem, 0, len(jobList)) + for _, j := range jobList { + label := j.Settings.Name + id := strconv.FormatInt(j.JobId, 10) + if label == "" { + label = id + } + out = append(out, ListItem{ID: id, Label: label}) + } + return out, nil +} + +// ListServingEndpoints lists serving endpoints as a paged result. +func ListServingEndpoints(ctx context.Context) (*PagedFetcher, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + iter := w.ServingEndpoints.List(ctx) + mapFn := func(e serving.ServingEndpoint) ListItem { + name := e.Name + if name == "" { + name = e.Id + } + return ListItem{ID: e.Id, Label: name} + } + items, hasMore, err := collectN(ctx, iter, pageSize, mapFn) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return collectN(ctx, iter, pageSize, mapFn) + }, + }, nil +} + +// ListExperiments lists MLflow experiments as a paged result. +func ListExperiments(ctx context.Context) (*PagedFetcher, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + iter := w.Experiments.ListExperiments(ctx, ml.ListExperimentsRequest{MaxResults: int64(pageSize)}) + mapFn := func(e ml.Experiment) ListItem { + label := e.Name + if label == "" { + label = e.ExperimentId + } + return ListItem{ID: e.ExperimentId, Label: label} + } + items, hasMore, err := collectN(ctx, iter, pageSize, mapFn) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return collectN(ctx, iter, pageSize, mapFn) + }, + }, nil +} + +// ListGenieSpaces lists Genie spaces as a paged result (manual pagination). +func ListGenieSpaces(ctx context.Context) (*PagedFetcher, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + var nextToken string + fetchPage := func(ctx context.Context) ([]ListItem, bool, error) { + req := dashboards.GenieListSpacesRequest{PageToken: nextToken, PageSize: pageSize} resp, err := w.Genie.ListSpaces(ctx, req) if err != nil { - return nil, err + return nil, false, err } + items := make([]ListItem, 0, len(resp.Spaces)) for _, s := range resp.Spaces { id := s.SpaceId label := s.Title @@ -463,54 +447,168 @@ func ListGenieSpaces(ctx context.Context) ([]ListItem, error) { if label == "" { label = id } - out = append(out, ListItem{ID: id, Label: label}) + items = append(items, ListItem{ID: id, Label: label}) } - if resp.NextPageToken == "" { - break + nextToken = resp.NextPageToken + return items, nextToken != "", nil + } + items, hasMore, err := fetchPage(ctx) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: fetchPage, + }, nil +} + +// ListConnections lists UC connections as a paged result. +func ListConnections(ctx context.Context) (*PagedFetcher, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + iter := w.Connections.List(ctx, catalog.ListConnectionsRequest{MaxResults: pageSize}) + mapFn := func(c catalog.ConnectionInfo) ListItem { + name := c.Name + if name == "" { + name = c.FullName } - req.PageToken = resp.NextPageToken + return ListItem{ID: c.FullName, Label: name} } - return out, nil + items, hasMore, err := collectN(ctx, iter, pageSize, mapFn) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return collectN(ctx, iter, pageSize, mapFn) + }, + }, nil } -// ListExperiments returns MLflow experiments as selectable items. -func ListExperiments(ctx context.Context) ([]ListItem, error) { +// ListVectorSearchIndexes lists vector search indexes as a paged result. +// Unlike other listers, this eagerly loads all indexes across all endpoints +// because the API requires a two-level query (endpoint -> indexes). +// Incremental loading isn't feasible without restructuring the picker into a +// two-step flow. The result is capped at maxTotalResults. +func ListVectorSearchIndexes(ctx context.Context) (*PagedFetcher, error) { w, err := workspaceClient(ctx) if err != nil { return nil, err } - iter := w.Experiments.ListExperiments(ctx, ml.ListExperimentsRequest{}) - exps, err := listing.ToSlice(ctx, iter) + var out []ListItem + epIter := w.VectorSearchEndpoints.ListEndpoints(ctx, vectorsearch.ListEndpointsRequest{}) + endpoints, err := listing.ToSlice(ctx, epIter) if err != nil { return nil, err } - out := make([]ListItem, 0, len(exps)) - for _, e := range exps { - label := e.Name - if label == "" { - label = e.ExperimentId + capped := false + for _, ep := range endpoints { + indexIter := w.VectorSearchIndexes.ListIndexes(ctx, vectorsearch.ListIndexesRequest{EndpointName: ep.Name}) + indexes, err := listing.ToSlice(ctx, indexIter) + if err != nil { + log.Warnf(ctx, "Failed to list indexes for endpoint %q: %v", ep.Name, err) + continue + } + for _, idx := range indexes { + label := idx.Name + if label == "" { + label = ep.Name + "/ (unnamed)" + } + id := ep.Name + "/" + idx.Name + out = append(out, ListItem{ID: id, Label: fmt.Sprintf("%s / %s", ep.Name, label)}) + if len(out) >= maxTotalResults { + capped = true + break + } + } + if capped { + break } - out = append(out, ListItem{ID: e.ExperimentId, Label: label}) } - return out, nil + return &PagedFetcher{Items: out, Capped: capped}, nil } -// TODO: uncomment when bundles support app as an app resource type. -// // ListAppsItems returns apps as ListItems (id = app name). -// func ListAppsItems(ctx context.Context) ([]ListItem, error) { -// w, err := workspaceClient(ctx) -// if err != nil { -// return nil, err -// } -// iter := w.Apps.List(ctx, apps.ListAppsRequest{}) -// appList, err := listing.ToSlice(ctx, iter) -// if err != nil { -// return nil, err -// } -// out := make([]ListItem, 0, len(appList)) -// for _, a := range appList { -// label := a.Name -// out = append(out, ListItem{ID: a.Name, Label: label}) -// } -// return out, nil -// } +// --------------------------------------------------------------------------- +// First-step paged constructors — used to prefetch the initial picker of +// multi-step resource prompts (catalog → schema → resource, etc.). +// --------------------------------------------------------------------------- + +// ListCatalogs lists UC catalogs as a paged result (first step of volume/function pickers). +func ListCatalogs(ctx context.Context) (*PagedFetcher, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + iter := w.Catalogs.List(ctx, catalog.ListCatalogsRequest{MaxResults: pageSize}) + mapFn := func(c catalog.CatalogInfo) ListItem { + return ListItem{ID: c.Name, Label: c.Name} + } + items, hasMore, err := collectN(ctx, iter, pageSize, mapFn) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return collectN(ctx, iter, pageSize, mapFn) + }, + }, nil +} + +// ListDatabaseInstances lists Lakebase database instances as a paged result. +func ListDatabaseInstances(ctx context.Context) (*PagedFetcher, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + apiLimit := min(pageSize, 100) + iter := w.Database.ListDatabaseInstances(ctx, database.ListDatabaseInstancesRequest{PageSize: apiLimit}) + mapFn := func(inst database.DatabaseInstance) ListItem { + return ListItem{ID: inst.Name, Label: inst.Name} + } + items, hasMore, err := collectN(ctx, iter, apiLimit, mapFn) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return collectN(ctx, iter, apiLimit, mapFn) + }, + }, nil +} + +// ListPostgresProjects lists Lakebase Autoscaling projects as a paged result. +func ListPostgresProjects(ctx context.Context) (*PagedFetcher, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + apiLimit := min(pageSize, 100) + iter := w.Postgres.ListProjects(ctx, postgres.ListProjectsRequest{PageSize: apiLimit}) + mapFn := func(p postgres.Project) ListItem { + label := p.Name + if p.Status != nil && p.Status.DisplayName != "" { + label = p.Status.DisplayName + } + return ListItem{ID: p.Name, Label: label} + } + items, hasMore, err := collectN(ctx, iter, apiLimit, mapFn) + if err != nil { + return nil, err + } + return &PagedFetcher{ + Items: items, + HasMore: hasMore, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return collectN(ctx, iter, apiLimit, mapFn) + }, + }, nil +} diff --git a/libs/apps/prompt/paged.go b/libs/apps/prompt/paged.go new file mode 100644 index 0000000000..97300e1f01 --- /dev/null +++ b/libs/apps/prompt/paged.go @@ -0,0 +1,98 @@ +package prompt + +import ( + "context" + + "github.com/databricks/databricks-sdk-go/listing" +) + +const ( + pageSize = 200 + maxTotalResults = 10000 + moreID = "__more__" + manualID = "__manual__" +) + +// PagedFetcher provides incremental access to a resource list. The first page +// is loaded in a background goroutine (signaled via the done channel). +// Subsequent pages are loaded on demand via LoadMore. Once maxTotalResults +// items have been accumulated, Capped is set to true and no more pages are +// offered — only the manual input fallback. +// +// Thread-safety: When created via startPrefetch, the fields are written by a +// background goroutine and the done channel is closed after all writes +// complete. Callers MUST call WaitForFirstPage (or verify IsDone returns +// true) before reading Items, HasMore, Capped, or Err. After the first page +// is ready, LoadMore must only be called from a single goroutine (the main +// prompt loop) — the underlying SDK iterator is not safe for concurrent use. +type PagedFetcher struct { + Items []ListItem + HasMore bool + Capped bool + Err error + + done chan struct{} // closed when the first page is ready + loadMore func(ctx context.Context) ([]ListItem, bool, error) +} + +// WaitForFirstPage blocks until the first page is ready or the context is cancelled. +func (p *PagedFetcher) WaitForFirstPage(ctx context.Context) error { + if p.done == nil { + return p.Err + } + select { + case <-p.done: + return p.Err + case <-ctx.Done(): + return ctx.Err() + } +} + +// IsDone returns true if the first page has already been loaded. +func (p *PagedFetcher) IsDone() bool { + if p.done == nil { + return true + } + select { + case <-p.done: + return true + default: + return false + } +} + +// LoadMore fetches the next page and appends it to Items. If the total reaches +// maxTotalResults, HasMore is cleared and Capped is set. +func (p *PagedFetcher) LoadMore(ctx context.Context) error { + if !p.HasMore || p.loadMore == nil { + return nil + } + items, hasMore, err := p.loadMore(ctx) + if err != nil { + return err + } + p.Items = append(p.Items, items...) + p.HasMore = hasMore + if len(p.Items) >= maxTotalResults { + p.HasMore = false + p.Capped = true + } + return nil +} + +// collectN consumes up to n items from an SDK iterator, mapping each to a +// ListItem. Returns the items, whether more exist, and any error. +func collectN[T any](ctx context.Context, iter listing.Iterator[T], n int, mapFn func(T) ListItem) ([]ListItem, bool, error) { + var items []ListItem + for len(items) < n { + if !iter.HasNext(ctx) { + return items, false, nil + } + item, err := iter.Next(ctx) + if err != nil { + return items, false, err + } + items = append(items, mapFn(item)) + } + return items, iter.HasNext(ctx), nil +} diff --git a/libs/apps/prompt/paged_test.go b/libs/apps/prompt/paged_test.go new file mode 100644 index 0000000000..ecfc6aa8fc --- /dev/null +++ b/libs/apps/prompt/paged_test.go @@ -0,0 +1,177 @@ +package prompt + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPagedFetcherIsDoneNilChannel(t *testing.T) { + f := &PagedFetcher{} + assert.True(t, f.IsDone()) +} + +func TestPagedFetcherIsDoneBeforeClose(t *testing.T) { + f := &PagedFetcher{done: make(chan struct{})} + assert.False(t, f.IsDone()) +} + +func TestPagedFetcherIsDoneAfterClose(t *testing.T) { + f := &PagedFetcher{done: make(chan struct{})} + close(f.done) + assert.True(t, f.IsDone()) +} + +func TestPagedFetcherWaitForFirstPageNilChannel(t *testing.T) { + f := &PagedFetcher{Err: errors.New("failed")} + err := f.WaitForFirstPage(t.Context()) + assert.EqualError(t, err, "failed") +} + +func TestPagedFetcherWaitForFirstPageSuccess(t *testing.T) { + f := &PagedFetcher{done: make(chan struct{})} + go func() { + f.Items = []ListItem{{ID: "1"}} + close(f.done) + }() + err := f.WaitForFirstPage(t.Context()) + assert.NoError(t, err) + assert.Len(t, f.Items, 1) +} + +func TestPagedFetcherWaitForFirstPageContextCancelled(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + f := &PagedFetcher{done: make(chan struct{})} + cancel() + + err := f.WaitForFirstPage(ctx) + assert.ErrorIs(t, err, context.Canceled) +} + +func TestPagedFetcherWaitForFirstPageError(t *testing.T) { + f := &PagedFetcher{done: make(chan struct{}), Err: errors.New("api error")} + close(f.done) + err := f.WaitForFirstPage(t.Context()) + assert.EqualError(t, err, "api error") +} + +func TestPagedFetcherLoadMoreAppendsItems(t *testing.T) { + f := &PagedFetcher{ + Items: []ListItem{{ID: "1"}}, + HasMore: true, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return []ListItem{{ID: "2"}, {ID: "3"}}, false, nil + }, + } + + err := f.LoadMore(t.Context()) + require.NoError(t, err) + assert.Len(t, f.Items, 3) + assert.False(t, f.HasMore) +} + +func TestPagedFetcherLoadMoreCapsAtLimit(t *testing.T) { + items := make([]ListItem, maxTotalResults-1) + for i := range items { + items[i] = ListItem{ID: "x"} + } + f := &PagedFetcher{ + Items: items, + HasMore: true, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return []ListItem{{ID: "last"}, {ID: "overflow"}}, true, nil + }, + } + + err := f.LoadMore(t.Context()) + require.NoError(t, err) + assert.False(t, f.HasMore) + assert.True(t, f.Capped) +} + +func TestPagedFetcherLoadMoreNoopWhenNotHasMore(t *testing.T) { + f := &PagedFetcher{ + Items: []ListItem{{ID: "1"}}, + HasMore: false, + } + err := f.LoadMore(t.Context()) + assert.NoError(t, err) + assert.Len(t, f.Items, 1) +} + +func TestPagedFetcherLoadMoreNoopWhenNilFunc(t *testing.T) { + f := &PagedFetcher{HasMore: true, loadMore: nil} + err := f.LoadMore(t.Context()) + assert.NoError(t, err) +} + +func TestPagedFetcherLoadMoreError(t *testing.T) { + f := &PagedFetcher{ + Items: []ListItem{{ID: "1"}}, + HasMore: true, + loadMore: func(ctx context.Context) ([]ListItem, bool, error) { + return nil, false, errors.New("network error") + }, + } + + err := f.LoadMore(t.Context()) + assert.EqualError(t, err, "network error") + assert.Len(t, f.Items, 1) +} + +// stubIterator implements listing.Iterator for testing collectN. +type stubIterator struct { + items []string + pos int +} + +func (s *stubIterator) HasNext(_ context.Context) bool { + return s.pos < len(s.items) +} + +func (s *stubIterator) Next(_ context.Context) (string, error) { + if s.pos >= len(s.items) { + return "", errors.New("exhausted") + } + item := s.items[s.pos] + s.pos++ + return item, nil +} + +func TestCollectNFullPage(t *testing.T) { + iter := &stubIterator{items: []string{"a", "b", "c", "d", "e"}} + items, hasMore, err := collectN(t.Context(), iter, 3, func(s string) ListItem { + return ListItem{ID: s, Label: s} + }) + + require.NoError(t, err) + assert.Len(t, items, 3) + assert.True(t, hasMore) + assert.Equal(t, "a", items[0].ID) + assert.Equal(t, "c", items[2].ID) +} + +func TestCollectNExhaustsIterator(t *testing.T) { + iter := &stubIterator{items: []string{"a", "b"}} + items, hasMore, err := collectN(t.Context(), iter, 5, func(s string) ListItem { + return ListItem{ID: s, Label: s} + }) + + require.NoError(t, err) + assert.Len(t, items, 2) + assert.False(t, hasMore) +} + +func TestCollectNEmptyIterator(t *testing.T) { + iter := &stubIterator{items: nil} + items, hasMore, err := collectN(t.Context(), iter, 5, func(s string) ListItem { + return ListItem{ID: s, Label: s} + }) + + require.NoError(t, err) + assert.Empty(t, items) + assert.False(t, hasMore) +} diff --git a/libs/apps/prompt/prefetch.go b/libs/apps/prompt/prefetch.go new file mode 100644 index 0000000000..d3f38bade0 --- /dev/null +++ b/libs/apps/prompt/prefetch.go @@ -0,0 +1,105 @@ +package prompt + +import ( + "context" + + "github.com/databricks/cli/libs/apps/manifest" + "github.com/databricks/cli/libs/log" +) + +// pagedConstructor creates a PagedFetcher with its first page loaded. +type pagedConstructor func(ctx context.Context) (*PagedFetcher, error) + +// pagedConstructors maps resource types to their paged lister constructor. +var pagedConstructors = map[string]pagedConstructor{ + ResourceTypeSQLWarehouse: ListSQLWarehouses, + ResourceTypeJob: ListJobs, + ResourceTypeServingEndpoint: ListServingEndpoints, + ResourceTypeGenieSpace: ListGenieSpaces, + ResourceTypeExperiment: ListExperiments, + ResourceTypeUCConnection: ListConnections, + ResourceTypeVectorSearchIndex: ListVectorSearchIndexes, +} + +// Internal cache keys for first-step fetchers used by multi-step prompts. +const ( + cacheKeyCatalogs = "_catalogs" + cacheKeyDatabaseInstances = "_database_instances" + cacheKeyPostgresProjects = "_postgres_projects" +) + +// firstStepPrefetch maps multi-step resource types to the cache key and +// constructor for their first picker step (e.g., volume → catalogs). +var firstStepPrefetch = map[string]struct { + cacheKey string + ctor pagedConstructor +}{ + ResourceTypeVolume: {cacheKeyCatalogs, ListCatalogs}, + ResourceTypeUCFunction: {cacheKeyCatalogs, ListCatalogs}, + ResourceTypeDatabase: {cacheKeyDatabaseInstances, ListDatabaseInstances}, + ResourceTypePostgres: {cacheKeyPostgresProjects, ListPostgresProjects}, +} + +// PrefetchResources kicks off a background goroutine for every resource type +// found in resources. Each goroutine fetches the first page (pageSize items) +// and stores a PagedFetcher with the iterator alive for subsequent LoadMore +// calls. The function returns immediately with a context carrying the cache. +// +// For single-step resources (warehouses, jobs, etc.) the fetcher is stored +// under the resource type key. For multi-step resources (volumes, functions, +// databases, postgres) the first-step fetcher (catalogs, instances, projects) +// is prefetched under an internal cache key so the first picker renders +// instantly. +// +// Goroutine lifecycle: each goroutine makes one SDK API call to fetch the +// first page. The goroutines respect the provided context — when the context +// is cancelled (e.g. Ctrl+C), the SDK HTTP client aborts the request and the +// goroutine terminates shortly after. Callers do not need to explicitly join +// the goroutines; they are short-lived and self-draining. +func PrefetchResources(ctx context.Context, resources []manifest.Resource) context.Context { + cache := CacheFromContext(ctx) + if cache == nil { + cache = NewResourceCache() + } + + for _, r := range resources { + // Single-step resources — full paged fetcher. + if ctor, ok := pagedConstructors[r.Type]; ok { + if cache.GetFetcher(r.Type) == nil { + log.Debugf(ctx, "Prefetching resource type %q", r.Type) + startPrefetch(ctx, cache, r.Type, ctor) + } + } + + // Multi-step resources — prefetch the first step. + if fs, ok := firstStepPrefetch[r.Type]; ok { + if cache.GetFetcher(fs.cacheKey) == nil { + log.Debugf(ctx, "Prefetching first step %q for resource type %q", fs.cacheKey, r.Type) + startPrefetch(ctx, cache, fs.cacheKey, fs.ctor) + } + } + } + + return ContextWithCache(ctx, cache) +} + +// startPrefetch launches a background goroutine that creates a PagedFetcher +// and stores it in the cache under the given key. The fetcher's done channel +// is closed after all fields are written, establishing a happens-before +// relationship — callers must call WaitForFirstPage before reading fields. +func startPrefetch(ctx context.Context, cache *ResourceCache, key string, ctor pagedConstructor) { + f := &PagedFetcher{done: make(chan struct{})} + cache.SetFetcher(key, f) + go func() { + defer close(f.done) + fetcher, err := ctor(ctx) + if err != nil { + f.Err = err + return + } + f.Items = fetcher.Items + f.HasMore = fetcher.HasMore + f.Capped = fetcher.Capped + f.loadMore = fetcher.loadMore + }() +} diff --git a/libs/apps/prompt/prompt.go b/libs/apps/prompt/prompt.go index ba448fa2a2..568140fa11 100644 --- a/libs/apps/prompt/prompt.go +++ b/libs/apps/prompt/prompt.go @@ -20,25 +20,27 @@ import ( "github.com/databricks/databricks-sdk-go/listing" "github.com/databricks/databricks-sdk-go/service/apps" "github.com/databricks/databricks-sdk-go/service/postgres" - "github.com/databricks/databricks-sdk-go/service/sql" ) // DefaultAppDescription is the default description for new apps. const DefaultAppDescription = "A Databricks App powered by AppKit" +// Brand palette — tuned for legibility on both light and dark terminals. +var ( + colorRed = lipgloss.Color("#E84040") // Bright Databricks red + colorGray = lipgloss.Color("#A1A1AA") // Light gray, legible on dark backgrounds + colorYellow = lipgloss.Color("#FFAB00") // Databricks yellow / amber + colorOrange = lipgloss.Color("#FF5F40") // Databricks orange (code blocks) +) + // AppkitTheme returns a custom theme for appkit prompts. func AppkitTheme() *huh.Theme { t := huh.ThemeBase() - // Databricks brand colors - red := lipgloss.Color("#BD2B26") - gray := lipgloss.Color("#71717A") // Mid-tone gray, readable on light and dark - yellow := lipgloss.Color("#FFAB00") - - t.Focused.Title = t.Focused.Title.Foreground(red).Bold(true) - t.Focused.Description = t.Focused.Description.Foreground(gray) - t.Focused.SelectedOption = t.Focused.SelectedOption.Foreground(yellow) - t.Focused.TextInput.Placeholder = t.Focused.TextInput.Placeholder.Foreground(gray) + t.Focused.Title = t.Focused.Title.Foreground(colorRed).Bold(true) + t.Focused.Description = t.Focused.Description.Foreground(colorGray) + t.Focused.SelectedOption = t.Focused.SelectedOption.Foreground(colorYellow) + t.Focused.TextInput.Placeholder = t.Focused.TextInput.Placeholder.Foreground(colorGray) return t } @@ -46,9 +48,9 @@ func AppkitTheme() *huh.Theme { // Styles for printing answered prompts. var ( answeredTitleStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color("#71717A")) + Foreground(colorGray) answeredValueStyle = lipgloss.NewStyle(). - Foreground(lipgloss.Color("#FFAB00")). + Foreground(colorYellow). Bold(true) ) @@ -118,11 +120,11 @@ func ValidateProjectName(s string) error { // PrintHeader prints the AppKit header banner. func PrintHeader(ctx context.Context) { headerStyle := lipgloss.NewStyle(). - Foreground(lipgloss.Color("#BD2B26")). + Foreground(colorRed). Bold(true) subtitleStyle := lipgloss.NewStyle(). - Foreground(lipgloss.Color("#71717A")) + Foreground(colorGray) cmdio.LogString(ctx, "") cmdio.LogString(ctx, headerStyle.Render("◆ Create a new Databricks AppKit project")) @@ -218,17 +220,6 @@ func PromptForDeployAndRun(ctx context.Context) (deploy bool, runMode RunMode, e return deploy, RunMode(runModeStr), nil } -// ListSQLWarehouses fetches all SQL warehouses the user has access to. -func ListSQLWarehouses(ctx context.Context) ([]sql.EndpointInfo, error) { - w := cmdctx.WorkspaceClient(ctx) - if w == nil { - return nil, errors.New("no workspace client available") - } - - iter := w.Warehouses.List(ctx, sql.ListWarehousesRequest{}) - return listing.ToSlice(ctx, iter) -} - // PromptFromList shows a picker for items and returns the selected ID. // If required is false and items are empty, returns ("", nil). If required is true and items are empty, returns an error. func PromptFromList(ctx context.Context, title, emptyMessage string, items []ListItem, required bool) (string, error) { @@ -254,10 +245,9 @@ func promptFromListWithLabel(ctx context.Context, title, emptyMessage string, it var selected string err := huh.NewSelect[string](). Title(title). - Description(fmt.Sprintf("%d available — type to filter", len(items))). + Description(fmt.Sprintf("%d available — / to filter", len(items))). Options(options...). Value(&selected). - Filtering(true). Height(8). WithTheme(theme). Run() @@ -268,6 +258,201 @@ func promptFromListWithLabel(ctx context.Context, title, emptyMessage string, it return selected, labels[selected], nil } +// awaitFetcher waits for a background PagedFetcher's first page. If the data +// is already available it returns immediately; otherwise a spinner is shown. +func awaitFetcher(ctx context.Context, f *PagedFetcher, spinnerMsg string) error { + if f.IsDone() { + return f.Err + } + return RunWithSpinnerCtx(ctx, spinnerMsg, func() error { + return f.WaitForFirstPage(ctx) + }) +} + +// getFetcher returns a PagedFetcher from the cache, waiting for its first page. +// If the cache has no entry, it creates one synchronously using the paged +// constructor registered in pagedConstructors. +func getFetcher(ctx context.Context, resourceType, spinnerMsg string) (*PagedFetcher, error) { + ctor := pagedConstructors[resourceType] + return getFetcherByKey(ctx, resourceType, spinnerMsg, ctor) +} + +// getFetcherByKey returns a PagedFetcher from the cache under the given key, +// waiting for its first page. If the cache has no entry, it falls back to +// creating one synchronously using the provided constructor. +func getFetcherByKey(ctx context.Context, cacheKey, spinnerMsg string, fallbackCtor pagedConstructor) (*PagedFetcher, error) { + if cache := CacheFromContext(ctx); cache != nil { + if f := cache.GetFetcher(cacheKey); f != nil { + if err := awaitFetcher(ctx, f, spinnerMsg); err != nil { + return nil, err + } + return f, nil + } + } + if fallbackCtor == nil { + return nil, fmt.Errorf("no lister registered for cache key %q", cacheKey) + } + var f *PagedFetcher + err := RunWithSpinnerCtx(ctx, spinnerMsg, func() error { + var fetchErr error + f, fetchErr = fallbackCtor(ctx) + return fetchErr + }) + if err != nil { + return nil, err + } + return f, nil +} + +// promptManualInput shows a text input for the user to type a resource name/ID +// manually. prefetchedLabels provides tab-complete suggestions. +func promptManualInput(ctx context.Context, title string, prefetchedLabels []string) (string, error) { + theme := AppkitTheme() + var value string + err := huh.NewInput(). + Title(title). + Placeholder("Type a name or ID"). + Suggestions(prefetchedLabels). + Value(&value). + WithTheme(theme). + Run() + if err != nil { + return "", err + } + return strings.TrimSpace(value), nil +} + +// promptFromPagedFetcher shows a picker backed by a PagedFetcher. When more +// pages are available and the total is under maxTotalResults, a "Load more..." +// option is appended. Once capped (>= maxTotalResults), an "Enter name/ID +// manually..." option replaces it. +// SearchFunc performs a server-side search by name/query. When non-nil, the +// manual input fallback triggers a search instead of accepting raw input. +// This is currently supported by Jobs (name filter). Other resource types can +// pass nil until their APIs add server-side filtering support. +type SearchFunc func(ctx context.Context, query string) ([]ListItem, error) + +func promptFromPagedFetcher(ctx context.Context, title, emptyMessage string, fetcher *PagedFetcher, required bool, searchFn SearchFunc) (string, string, error) { + if len(fetcher.Items) == 0 && !fetcher.HasMore { + if required { + return "", "", errors.New(emptyMessage) + } + return "", "", nil + } + theme := AppkitTheme() + + for { + options := make([]huh.Option[string], 0, len(fetcher.Items)+2) + labels := make(map[string]string, len(fetcher.Items)) + + var desc string + if fetcher.HasMore && !fetcher.Capped { + desc = fmt.Sprintf("%d results loaded — / to search", len(fetcher.Items)) + options = append(options, huh.NewOption("+ Load more results", moreID)) + } else if fetcher.Capped { + desc = fmt.Sprintf("%d results loaded — / to search", len(fetcher.Items)) + manualLabel := "Not listed? Enter ID manually..." + if searchFn != nil { + manualLabel = "Not listed? Search by name..." + } + options = append(options, huh.NewOption(manualLabel, manualID)) + } else { + desc = fmt.Sprintf("%d available — / to search", len(fetcher.Items)) + } + + for _, it := range fetcher.Items { + options = append(options, huh.NewOption(it.Label, it.ID)) + labels[it.ID] = it.Label + } + + var selected string + err := huh.NewSelect[string](). + Title(title). + Description(desc). + Options(options...). + Value(&selected). + Height(8). + WithTheme(theme). + Run() + if err != nil { + return "", "", err + } + + switch selected { + case moreID: + if err := RunWithSpinnerCtx(ctx, "Fetching more results...", func() error { + return fetcher.LoadMore(ctx) + }); err != nil { + return "", "", err + } + continue + + case manualID: + suggestions := make([]string, 0, len(fetcher.Items)) + for _, it := range fetcher.Items { + suggestions = append(suggestions, it.Label) + } + query, inputErr := promptManualInput(ctx, title, suggestions) + if inputErr != nil { + return "", "", inputErr + } + if query == "" { + if required { + continue + } + return "", "", nil + } + + if searchFn == nil { + printAnswered(ctx, title, query) + return query, query, nil + } + + var results []ListItem + if searchErr := RunWithSpinnerCtx(ctx, fmt.Sprintf("Searching for %q...", query), func() error { + var fetchErr error + results, fetchErr = searchFn(ctx, query) + return fetchErr + }); searchErr != nil { + return "", "", searchErr + } + if len(results) == 0 { + printAnswered(ctx, title, query) + return query, query, nil + } + if len(results) == 1 { + printAnswered(ctx, title, results[0].Label) + return results[0].ID, results[0].Label, nil + } + id, label, pickErr := promptFromListWithLabel(ctx, title+" — search results", "no matches", results, required) + if pickErr != nil { + return "", "", pickErr + } + return id, label, nil + + default: + printAnswered(ctx, title, labels[selected]) + return selected, labels[selected], nil + } + } +} + +// promptForPagedResource gets a PagedFetcher (from cache or on-demand), then +// shows the paged picker with Load more / Enter manually support. +// Pass a non-nil searchFn to enable server-side search in the manual input +// fallback (currently only Jobs supports this). +func promptForPagedResource(ctx context.Context, r manifest.Resource, required bool, title, emptyMsg, spinnerMsg string, searchFn SearchFunc) (map[string]string, error) { + f, err := getFetcher(ctx, r.Type, spinnerMsg) + if err != nil { + return nil, err + } + value, _, promptErr := promptFromPagedFetcher(ctx, title, emptyMsg, f, required, searchFn) + if promptErr != nil { + return nil, promptErr + } + return singleValueResult(r, value), nil +} + // PromptForWarehouse shows a picker to select a SQL warehouse. func PromptForWarehouse(ctx context.Context) (string, error) { var items []ListItem @@ -282,6 +467,19 @@ func PromptForWarehouse(ctx context.Context) (string, error) { return PromptFromList(ctx, "Select SQL Warehouse", "no SQL warehouses found. Create one in your workspace first", items, true) } +// resourceTitle returns a prompt title for a resource, including the plugin name +// for context when available (e.g. "Select SQL Warehouse for Analytics"). +func resourceTitle(fallback string, r manifest.Resource) string { + title := r.Alias + if title == "" { + title = fallback + } + if r.PluginDisplayName != "" { + title = fmt.Sprintf("%s for %s", title, r.PluginDisplayName) + } + return title +} + // singleValueResult wraps a single value into the resource values map. // Uses the first field name from Fields for the composite key (resource_key.field), // or falls back to the resource key if no Fields are defined. @@ -296,24 +494,6 @@ func singleValueResult(r manifest.Resource, value string) map[string]string { return map[string]string{r.Key(): value} } -// promptForResourceFromLister runs a spinner, fetches items via fn, then shows PromptFromList. -func promptForResourceFromLister(ctx context.Context, r manifest.Resource, required bool, title, emptyMsg, spinnerMsg string, fn func(context.Context) ([]ListItem, error)) (map[string]string, error) { - var items []ListItem - err := RunWithSpinnerCtx(ctx, spinnerMsg, func() error { - var fetchErr error - items, fetchErr = fn(ctx) - return fetchErr - }) - if err != nil { - return nil, err - } - value, err := PromptFromList(ctx, title, emptyMsg, items, required) - if err != nil { - return nil, err - } - return singleValueResult(r, value), nil -} - // PromptForSecret shows a two-step picker for secret scope and key. func PromptForSecret(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { // Step 1: pick scope @@ -358,30 +538,30 @@ func PromptForSecret(ctx context.Context, r manifest.Resource, required bool) (m }, nil } -// PromptForJob shows a picker for jobs. +// PromptForJob shows a picker for jobs. When the user selects "Enter manually" +// (after the 500-item cap), the input triggers a server-side name search via +// the Jobs API's Name filter before accepting the value. func PromptForJob(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - return promptForResourceFromLister(ctx, r, required, "Select Job", "no jobs found", "Fetching jobs...", ListJobs) + title := resourceTitle("Select Job", r) + return promptForPagedResource(ctx, r, required, title, "no jobs found", "Fetching jobs...", SearchJobs) } // PromptForSQLWarehouseResource shows a picker for SQL warehouses (manifest.Resource version). func PromptForSQLWarehouseResource(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - return promptForResourceFromLister(ctx, r, required, "Select SQL Warehouse", "no SQL warehouses found. Create one in your workspace first", "Fetching SQL warehouses...", ListSQLWarehousesItems) + title := resourceTitle("Select SQL Warehouse", r) + return promptForPagedResource(ctx, r, required, title, "no SQL warehouses found. Create one in your workspace first", "Fetching SQL warehouses...", nil) } const backID = "__back__" // promptUCCatalog shows a picker for UC catalogs (shared first step for volume/function pickers). func promptUCCatalog(ctx context.Context, required bool) (string, error) { - var items []ListItem - err := RunWithSpinnerCtx(ctx, "Fetching catalogs...", func() error { - var fetchErr error - items, fetchErr = ListCatalogs(ctx) - return fetchErr - }) + f, err := getFetcherByKey(ctx, cacheKeyCatalogs, "Fetching catalogs...", ListCatalogs) if err != nil { return "", err } - return PromptFromList(ctx, "Select Catalog", "no catalogs found", items, required) + id, _, promptErr := promptFromPagedFetcher(ctx, "Select Catalog", "no catalogs found", f, required, nil) + return id, promptErr } // promptFromListWithBack shows a picker with a "← Go back" option prepended. @@ -460,17 +640,45 @@ func promptUCResource(ctx context.Context, r manifest.Resource, required bool, r // PromptForServingEndpoint shows a picker for serving endpoints. func PromptForServingEndpoint(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - return promptForResourceFromLister(ctx, r, required, "Select Serving Endpoint", "no serving endpoints found", "Fetching serving endpoints...", ListServingEndpoints) + title := resourceTitle("Select Serving Endpoint", r) + return promptForPagedResource(ctx, r, required, title, "no serving endpoints found", "Fetching serving endpoints...", nil) +} + +// volumePathToSecurableName converts a volume path (/Volumes/catalog/schema/vol) +// to the securable_full_name format (catalog.schema.vol) used by DABs. +func volumePathToSecurableName(path string) string { + trimmed := strings.TrimPrefix(path, "/Volumes/") + if trimmed == path { + return path + } + return strings.ReplaceAll(trimmed, "/", ".") } // PromptForVolume shows a three-step picker for UC volumes: catalog -> schema -> volume. +// Stores two values: the volume path for .env and the dot-separated securable +// name for the DABs YAML securable_full_name field. func PromptForVolume(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - return promptUCResource(ctx, r, required, "Volume", "Fetching volumes...", ListVolumesInSchema) + result, err := promptUCResource(ctx, r, required, "Volume", "Fetching volumes...", ListVolumesInSchema) + if err != nil || result == nil { + return result, err + } + // promptUCResource stores the volume path (/Volumes/cat/schema/vol) under + // the first manifest field (e.g., "path"). The DABs spec also needs the + // securable_full_name (cat.schema.vol) under "id". + idKey := r.Key() + ".id" + if _, exists := result[idKey]; !exists { + for _, v := range result { + result[idKey] = volumePathToSecurableName(v) + break + } + } + return result, nil } // PromptForVectorSearchIndex shows a picker for vector search indexes. func PromptForVectorSearchIndex(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - return promptForResourceFromLister(ctx, r, required, "Select Vector Search Index", "no vector search indexes found", "Fetching vector search indexes...", ListVectorSearchIndexes) + title := resourceTitle("Select Vector Search Index", r) + return promptForPagedResource(ctx, r, required, title, "no vector search indexes found", "Fetching vector search indexes...", nil) } // PromptForUCFunction shows a three-step picker for UC functions: catalog -> schema -> function. @@ -480,22 +688,18 @@ func PromptForUCFunction(ctx context.Context, r manifest.Resource, required bool // PromptForUCConnection shows a picker for UC connections. func PromptForUCConnection(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - return promptForResourceFromLister(ctx, r, required, "Select UC Connection", "no connections found", "Fetching connections...", ListConnections) + title := resourceTitle("Select UC Connection", r) + return promptForPagedResource(ctx, r, required, title, "no connections found", "Fetching connections...", nil) } // PromptForDatabase shows a two-step picker for database instance and database name. func PromptForDatabase(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { // Step 1: pick a Lakebase instance - var instances []ListItem - err := RunWithSpinnerCtx(ctx, "Fetching database instances...", func() error { - var fetchErr error - instances, fetchErr = ListDatabaseInstances(ctx) - return fetchErr - }) + f, err := getFetcherByKey(ctx, cacheKeyDatabaseInstances, "Fetching database instances...", ListDatabaseInstances) if err != nil { return nil, err } - instanceName, err := PromptFromList(ctx, "Select Database Instance", "no database instances found", instances, required) + instanceName, _, err := promptFromPagedFetcher(ctx, "Select Database Instance", "no database instances found", f, required, nil) if err != nil { return nil, err } @@ -530,16 +734,11 @@ func PromptForDatabase(ctx context.Context, r manifest.Resource, required bool) // PromptForPostgres shows a three-step picker for Lakebase Autoscaling (V2): project, branch, then database. func PromptForPostgres(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { // Step 1: pick a project - var projects []ListItem - err := RunWithSpinnerCtx(ctx, "Fetching Postgres projects...", func() error { - var fetchErr error - projects, fetchErr = ListPostgresProjects(ctx) - return fetchErr - }) + f, err := getFetcherByKey(ctx, cacheKeyPostgresProjects, "Fetching Postgres projects...", ListPostgresProjects) if err != nil { return nil, err } - projectName, err := PromptFromList(ctx, "Select Postgres Project", "no Postgres projects found", projects, required) + projectName, _, err := promptFromPagedFetcher(ctx, "Select Postgres Project", "no Postgres projects found", f, required, nil) if err != nil { return nil, err } @@ -674,16 +873,13 @@ func applyResolvedValues(r manifest.Resource, resolvedValues, result map[string] // PromptForGenieSpace shows a picker for Genie spaces. // Captures both the space ID and name since the DABs schema requires both fields. func PromptForGenieSpace(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - var items []ListItem - err := RunWithSpinnerCtx(ctx, "Fetching Genie spaces...", func() error { - var fetchErr error - items, fetchErr = ListGenieSpaces(ctx) - return fetchErr - }) + f, err := getFetcher(ctx, r.Type, "Fetching Genie spaces...") if err != nil { return nil, err } - id, name, err := promptFromListWithLabel(ctx, "Select Genie Space", "no Genie spaces found", items, required) + + title := resourceTitle("Select Genie Space", r) + id, name, err := promptFromPagedFetcher(ctx, title, "no Genie spaces found", f, required, nil) if err != nil { return nil, err } @@ -698,7 +894,8 @@ func PromptForGenieSpace(ctx context.Context, r manifest.Resource, required bool // PromptForExperiment shows a picker for MLflow experiments. func PromptForExperiment(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - return promptForResourceFromLister(ctx, r, required, "Select Experiment", "no experiments found", "Fetching experiments...", ListExperiments) + title := resourceTitle("Select Experiment", r) + return promptForPagedResource(ctx, r, required, title, "no experiments found", "Fetching experiments...", nil) } // TODO: uncomment when bundles support app as an app resource type. @@ -707,7 +904,27 @@ func PromptForExperiment(ctx context.Context, r manifest.Resource, required bool // return promptForResourceFromLister(ctx, r, required, "Select App", "no apps found. Create one first with 'databricks apps create '", "Fetching apps...", ListAppsItems) // } +// Styles for consistent status output. +var ( + doneStyle = lipgloss.NewStyle(). + Foreground(colorYellow). + Bold(true) + doneTextStyle = lipgloss.NewStyle(). + Foreground(colorGray) +) + +// PrintDone prints a styled "✔ message" completion line. +func PrintDone(ctx context.Context, msg string) { + cmdio.LogString(ctx, fmt.Sprintf("%s %s", doneStyle.Render("✔"), doneTextStyle.Render(msg))) +} + +// stripEllipsis removes a trailing "..." from a string for use in completion messages. +func stripEllipsis(s string) string { + return strings.TrimSuffix(s, "...") +} + // RunWithSpinnerCtx runs a function while showing a spinner with the given title. +// On success, prints a styled checkmark completion line. // The spinner stops and the function returns early if the context is cancelled. // Panics in the action are recovered and returned as errors. func RunWithSpinnerCtx(ctx context.Context, title string, action func() error) error { @@ -727,6 +944,9 @@ func RunWithSpinnerCtx(ctx context.Context, title string, action func() error) e select { case err := <-done: spinner.Close() + if err == nil { + PrintDone(ctx, stripEllipsis(title)) + } return err case <-ctx.Done(): spinner.Close() @@ -789,10 +1009,9 @@ func PromptForAppSelection(ctx context.Context, title string) (string, error) { var selected string err = huh.NewSelect[string](). Title(title). - Description(fmt.Sprintf("%d apps found — type to filter", len(existingApps))). + Description(fmt.Sprintf("%d apps found — / to filter", len(existingApps))). Options(options...). Value(&selected). - Filtering(true). Height(8). WithTheme(theme). Run() @@ -808,14 +1027,14 @@ func PromptForAppSelection(ctx context.Context, title string) (string, error) { // If nextStepsCmd is non-empty, also prints the "Next steps" section with the given command. func PrintSuccess(ctx context.Context, projectName, outputDir string, fileCount int, nextStepsCmd string) { successStyle := lipgloss.NewStyle(). - Foreground(lipgloss.Color("#FFAB00")). // Databricks yellow + Foreground(colorYellow). Bold(true) dimStyle := lipgloss.NewStyle(). - Foreground(lipgloss.Color("#71717A")) // Mid-tone gray + Foreground(colorGray) codeStyle := lipgloss.NewStyle(). - Foreground(lipgloss.Color("#FF3621")) // Databricks orange + Foreground(colorOrange) cmdio.LogString(ctx, "") cmdio.LogString(ctx, successStyle.Render("✔ Project created successfully!")) @@ -842,15 +1061,15 @@ type SetupNote struct { // PrintSetupNotes renders a styled "Setup Notes" section for selected plugins. func PrintSetupNotes(ctx context.Context, notes []SetupNote) { headerStyle := lipgloss.NewStyle(). - Foreground(lipgloss.Color("#FFAB00")). // Databricks yellow + Foreground(colorYellow). Bold(true) nameStyle := lipgloss.NewStyle(). - Foreground(lipgloss.Color("#71717A")). // Mid-tone gray + Foreground(colorGray). Bold(true) msgStyle := lipgloss.NewStyle(). - Foreground(lipgloss.Color("#71717A")) // Mid-tone gray + Foreground(colorGray) cmdio.LogString(ctx, headerStyle.Render(" Setup Notes")) cmdio.LogString(ctx, "") diff --git a/libs/apps/prompt/prompt_test.go b/libs/apps/prompt/prompt_test.go index 794f0f6b9b..3400eab9be 100644 --- a/libs/apps/prompt/prompt_test.go +++ b/libs/apps/prompt/prompt_test.go @@ -274,6 +274,23 @@ func TestApplyResolvedValues(t *testing.T) { }) } +func TestVolumePathToSecurableName(t *testing.T) { + tests := []struct { + input string + want string + }{ + {"/Volumes/catalog/schema/vol", "catalog.schema.vol"}, + {"/Volumes/my-cat/my-schema/my-vol", "my-cat.my-schema.my-vol"}, + {"catalog.schema.vol", "catalog.schema.vol"}, + {"/Volumes/a/b/c", "a.b.c"}, + } + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + assert.Equal(t, tt.want, volumePathToSecurableName(tt.input)) + }) + } +} + func TestMaxAppNameLength(t *testing.T) { // Verify the constant is set correctly assert.Equal(t, 30, MaxAppNameLength)