package project import ( "bytes" "context" "errors" "fmt" "os" "strings" "github.com/databricks/cli/cmd/labs/github" "github.com/databricks/cli/cmd/labs/unpack" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg/cfgpickers" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/log" "github.com/databricks/cli/libs/process" "github.com/databricks/cli/libs/python" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/compute" "github.com/databricks/databricks-sdk-go/service/sql" "github.com/fatih/color" "github.com/spf13/cobra" ) const ownerRWXworldRX = 0o755 type whTypes []sql.EndpointInfoWarehouseType type hook struct { *Entrypoint `yaml:",inline"` Script string `yaml:"script"` RequireDatabricksConnect bool `yaml:"require_databricks_connect,omitempty"` MinRuntimeVersion string `yaml:"min_runtime_version,omitempty"` WarehouseTypes whTypes `yaml:"warehouse_types,omitempty"` } func (h *hook) RequireRunningCluster() bool { if h.Entrypoint == nil { return false } return h.Entrypoint.RequireRunningCluster } func (h *hook) HasPython() bool { return strings.HasSuffix(h.Script, ".py") } func (h *hook) runHook(cmd *cobra.Command) error { if h.Script == "" { return nil } ctx := cmd.Context() envs, err := h.Prepare(cmd) if err != nil { return fmt.Errorf("prepare: %w", err) } libDir := h.EffectiveLibDir() args := []string{} if strings.HasSuffix(h.Script, ".py") { args = append(args, h.virtualEnvPython(ctx)) } return process.Forwarded(ctx, append(args, h.Script), cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr(), process.WithDir(libDir), process.WithEnvs(envs)) } type installer struct { *Project version string // command instance is used for: // - auth profile flag override // - standard input, output, and error streams cmd *cobra.Command } func (i *installer) Install(ctx context.Context) error { err := i.EnsureFoldersExist() if err != nil { return fmt.Errorf("folders: %w", err) } i.folder, err = PathInLabs(ctx, i.Name) if err != nil { return err } w, err := i.login(ctx) if err != nil && errors.Is(err, profile.ErrNoConfiguration) { cfg, err := i.Installer.envAwareConfig(ctx) if err != nil { return err } w, err = databricks.NewWorkspaceClient((*databricks.Config)(cfg)) if err != nil { return fmt.Errorf("no ~/.databrickscfg: %w", err) } } else if err != nil { return fmt.Errorf("login: %w", err) } err = i.downloadLibrary(ctx) if err != nil { return fmt.Errorf("lib: %w", err) } err = i.setupPythonVirtualEnvironment(ctx, w) if err != nil { return fmt.Errorf("python: %w", err) } err = i.recordVersion(ctx) if err != nil { return fmt.Errorf("record version: %w", err) } // TODO: failing install hook for "clean installations" (not upgrages) // should trigger removal of the project, otherwise users end up with // misconfigured CLIs err = i.runInstallHook(ctx) if err != nil { return fmt.Errorf("installer: %w", err) } return nil } func (i *installer) Upgrade(ctx context.Context) error { err := i.downloadLibrary(ctx) if err != nil { return fmt.Errorf("lib: %w", err) } err = i.recordVersion(ctx) if err != nil { return fmt.Errorf("record version: %w", err) } err = i.runInstallHook(ctx) if err != nil { return fmt.Errorf("installer: %w", err) } err = i.installPythonDependencies(ctx, ".") if err != nil { return fmt.Errorf("python dependencies: %w", err) } return nil } func (i *installer) warningf(text string, v ...any) { i.cmd.PrintErrln(color.YellowString(text, v...)) } func (i *installer) cleanupLib(ctx context.Context) error { libDir := i.LibDir() err := os.RemoveAll(libDir) if err != nil { return fmt.Errorf("remove all: %w", err) } return os.MkdirAll(libDir, ownerRWXworldRX) } func (i *installer) recordVersion(ctx context.Context) error { return i.writeVersionFile(ctx, i.version) } func (i *installer) login(ctx context.Context) (*databricks.WorkspaceClient, error) { if !cmdio.IsPromptSupported(ctx) { log.Debugf(ctx, "Skipping workspace profile prompts in non-interactive mode") return nil, nil } cfg, err := i.metaEntrypoint(ctx).validLogin(i.cmd) if errors.Is(err, ErrNoLoginConfig) { cfg, err = i.Installer.envAwareConfig(ctx) if err != nil { return nil, err } } else if err != nil { return nil, fmt.Errorf("valid: %w", err) } if !i.HasAccountLevelCommands() && cfg.IsAccountClient() { return nil, fmt.Errorf("got account-level client, but no account-level commands") } lc := &loginConfig{Entrypoint: i.Installer.Entrypoint} w, err := lc.askWorkspace(ctx, cfg) if err != nil { return nil, fmt.Errorf("ask for workspace: %w", err) } err = lc.save(ctx) if err != nil { return nil, fmt.Errorf("save: %w", err) } return w, nil } func (i *installer) downloadLibrary(ctx context.Context) error { feedback := cmdio.Spinner(ctx) defer close(feedback) feedback <- "Cleaning up previous installation if necessary" err := i.cleanupLib(ctx) if err != nil { return fmt.Errorf("cleanup: %w", err) } libTarget := i.LibDir() // we may support wheels, jars, and golang binaries. but those are not zipballs if i.IsZipball() { feedback <- fmt.Sprintf("Downloading and unpacking zipball for %s", i.version) return i.downloadAndUnpackZipball(ctx, libTarget) } return fmt.Errorf("we only support zipballs for now") } func (i *installer) downloadAndUnpackZipball(ctx context.Context, libTarget string) error { raw, err := github.DownloadZipball(ctx, "databrickslabs", i.Name, i.version) if err != nil { return fmt.Errorf("download zipball from GitHub: %w", err) } zipball := unpack.GitHubZipball{Reader: bytes.NewBuffer(raw)} log.Debugf(ctx, "Unpacking zipball to: %s", libTarget) return zipball.UnpackTo(libTarget) } func (i *installer) setupPythonVirtualEnvironment(ctx context.Context, w *databricks.WorkspaceClient) error { if !i.HasPython() { return nil } feedback := cmdio.Spinner(ctx) defer close(feedback) feedback <- "Detecting all installed Python interpreters on the system" pythonInterpreters, err := python.DetectInterpreters(ctx) if err != nil { return fmt.Errorf("detect: %w", err) } py, err := pythonInterpreters.AtLeast(i.MinPython) if err != nil { return fmt.Errorf("min version: %w", err) } log.Debugf(ctx, "Detected Python %s at: %s", py.Version, py.Path) venvPath := i.virtualEnvPath(ctx) log.Debugf(ctx, "Creating Python Virtual Environment at: %s", venvPath) feedback <- fmt.Sprintf("Creating Virtual Environment with Python %s", py.Version) _, err = process.Background(ctx, []string{py.Path, "-m", "venv", venvPath}) if err != nil { return fmt.Errorf("create venv: %w", err) } if i.Installer != nil && i.Installer.RequireDatabricksConnect { feedback <- "Determining Databricks Connect version" cluster, err := w.Clusters.Get(ctx, compute.GetClusterRequest{ ClusterId: w.Config.ClusterID, }) if err != nil { return fmt.Errorf("cluster: %w", err) } runtimeVersion, ok := cfgpickers.GetRuntimeVersion(*cluster) if !ok { return fmt.Errorf("unsupported runtime: %s", cluster.SparkVersion) } feedback <- fmt.Sprintf("Installing Databricks Connect v%s", runtimeVersion) pipSpec := fmt.Sprintf("databricks-connect==%s", runtimeVersion) err = i.installPythonDependencies(ctx, pipSpec) if err != nil { return fmt.Errorf("dbconnect: %w", err) } } feedback <- "Installing Python library dependencies" return i.installPythonDependencies(ctx, ".") } func (i *installer) installPythonDependencies(ctx context.Context, spec string) error { if !i.IsPythonProject() { return nil } libDir := i.LibDir() log.Debugf(ctx, "Installing Python dependencies for: %s", libDir) // maybe we'll need to add call one of the two scripts: // - python3 -m ensurepip --default-pip // - curl -o https://bootstrap.pypa.io/get-pip.py | python3 var buf bytes.Buffer _, err := process.Background(ctx, []string{i.virtualEnvPython(ctx), "-m", "pip", "install", spec}, process.WithCombinedOutput(&buf), process.WithDir(libDir)) if err != nil { i.warningf(buf.String()) return fmt.Errorf("failed to install dependencies of %s", spec) } return nil } func (i *installer) runInstallHook(ctx context.Context) error { if i.Installer == nil { return nil } if i.Installer.Script == "" { return nil } log.Debugf(ctx, "Launching installer script %s in %s", i.Installer.Script, i.LibDir()) return i.Installer.runHook(i.cmd) }