databricks-cli/project/project.go

272 lines
7.4 KiB
Go
Raw Normal View History

2022-05-13 13:30:22 +00:00
package project
import (
"context"
"fmt"
"os"
"path/filepath"
2022-05-14 17:55:00 +00:00
"sync"
2022-05-13 13:30:22 +00:00
"github.com/databricks/bricks/libs/git"
"github.com/databricks/databricks-sdk-go"
"github.com/databricks/databricks-sdk-go/service/commands"
"github.com/databricks/databricks-sdk-go/service/scim"
"github.com/spf13/cobra"
2022-05-13 13:30:22 +00:00
)
const CacheDirName = ".databricks"
type project struct {
mu sync.Mutex
2022-05-14 17:55:00 +00:00
root string
env string
config *Config
environment *Environment
wsc *databricks.WorkspaceClient
me *scim.User
fileSet *git.FileSet
}
// Configure is used as a PreRunE function for all commands that
// require a project to be configured. If a project could successfully
// be found and loaded, it is set on the command's context object.
func Configure(cmd *cobra.Command, args []string) error {
root, err := getRoot()
if err != nil {
return err
}
2022-05-14 17:55:00 +00:00
ctx, err := Initialize(cmd.Context(), root, getEnvironment(cmd))
if err != nil {
return err
}
cmd.SetContext(ctx)
return nil
2022-05-14 17:55:00 +00:00
}
// Placeholder to use as unique key in context.Context.
var projectKey int
// Initialize loads a project configuration given a root and environment.
// It stores the project on a new context.
// The project is available through the `Get()` function.
func Initialize(ctx context.Context, root, env string) (context.Context, error) {
config, err := loadProjectConf(root)
if err != nil {
return nil, err
}
// Confirm that the specified environment is valid.
environment, ok := config.Environments[env]
if !ok {
return nil, fmt.Errorf("environment [%s] not defined", env)
}
fileSet, err := git.NewFileSet(root)
if err != nil {
return nil, err
}
err = fileSet.EnsureValidGitIgnoreExists()
if err != nil {
return nil, err
}
p := project{
root: root,
env: env,
config: &config,
environment: &environment,
fileSet: fileSet,
}
p.initializeWorkspacesClient(ctx)
return context.WithValue(ctx, &projectKey, &p), nil
2022-05-14 17:55:00 +00:00
}
func (p *project) initializeWorkspacesClient(ctx context.Context) {
var config databricks.Config
// If the config specifies a profile, or other authentication related properties,
// pass them along to the SDK here. If nothing is defined, the SDK will figure
// out which autentication mechanism to use using enviroment variables.
if p.environment.Workspace.Profile != "" {
config.Profile = p.environment.Workspace.Profile
}
p.wsc = databricks.Must(databricks.NewWorkspaceClient(&config))
}
// Get returns the project as configured on the context.
// It panics if it isn't configured.
func Get(ctx context.Context) *project {
project, ok := ctx.Value(&projectKey).(*project)
if !ok {
panic(`context not configured with project`)
}
return project
2022-05-13 13:30:22 +00:00
}
// Make sure to initialize the workspaces client on project init
func (p *project) WorkspacesClient() *databricks.WorkspaceClient {
return p.wsc
}
func (p *project) Root() string {
return p.root
}
func (p *project) GetFileSet() *git.FileSet {
return p.fileSet
}
// This cache dir will contain any state, state overrides (per user overrides
// to the project config) or any generated artifacts (eg: sync snapshots)
// that should never be checked into Git.
//
// We enfore that cache dir (.databricks) is added to .gitignore
// because it contains per-user overrides that we do not want users to
// accidentally check into git
func (p *project) CacheDir() (string, error) {
// assert cache dir is present in git ignore
ign, err := p.fileSet.IgnoreDirectory(fmt.Sprintf("/%s/", CacheDirName))
if err != nil {
return "", fmt.Errorf("failed to check if directory %s is ignored: %w", CacheDirName, err)
}
if !ign {
return "", fmt.Errorf("please add /%s/ to .gitignore", CacheDirName)
}
cacheDirPath := filepath.Join(p.root, CacheDirName)
// create cache dir if it does not exist
if _, err := os.Stat(cacheDirPath); os.IsNotExist(err) {
err = os.Mkdir(cacheDirPath, os.ModeDir|os.ModePerm)
if err != nil {
return "", fmt.Errorf("failed to create cache directory %s with error: %s", cacheDirPath, err)
}
}
return cacheDirPath, nil
}
func (p *project) Config() Config {
return *p.config
}
func (p *project) Environment() Environment {
return *p.environment
}
func (p *project) Me() (*scim.User, error) {
p.mu.Lock()
defer p.mu.Unlock()
if p.me != nil {
return p.me, nil
2022-05-14 17:55:00 +00:00
}
me, err := p.wsc.CurrentUser.Me(context.Background())
2022-05-13 13:30:22 +00:00
if err != nil {
return nil, err
2022-05-14 17:55:00 +00:00
}
p.me = me
return me, nil
2022-05-14 17:55:00 +00:00
}
func (p *project) DeploymentIsolationPrefix() string {
if p.config.Isolation == None {
return p.config.Name
2022-05-21 13:23:37 +00:00
}
if p.config.Isolation == Soft {
me, err := p.Me()
if err != nil {
panic(err)
}
return fmt.Sprintf("%s/%s", p.config.Name, me.UserName)
2022-05-21 13:23:37 +00:00
}
panic(fmt.Errorf("unknow project isolation: %s", p.config.Isolation))
2022-05-21 13:23:37 +00:00
}
func getClusterIdFromClusterName(ctx context.Context,
wsc *databricks.WorkspaceClient,
clusterName string,
) (clusterId string, err error) {
clusterInfo, err := wsc.Clusters.GetByClusterName(ctx, clusterName)
if err != nil {
return "", err
}
return clusterInfo.ClusterId, nil
}
// Old version of getting development cluster details with isolation implemented.
// Kept just for reference. Remove once isolation is implemented properly
/*
func (p *project) DevelopmentCluster(ctx context.Context) (cluster clusters.ClusterInfo, err error) {
api := clusters.NewClustersAPI(ctx, p.Client()) // TODO: rewrite with normal SDK
if p.project.DevCluster == nil {
p.project.DevCluster = &clusters.Cluster{}
}
dc := p.project.DevCluster
if p.project.Isolation == Soft {
if p.project.IsDevClusterJustReference() {
2022-05-14 17:55:00 +00:00
err = fmt.Errorf("projects with soft isolation cannot have named clusters")
return
}
dc.ClusterName = fmt.Sprintf("dev/%s", p.DeploymentIsolationPrefix())
2022-05-14 17:55:00 +00:00
}
if dc.ClusterName == "" {
err = fmt.Errorf("please either pick `isolation: soft` or specify a shared cluster name")
2022-05-13 13:30:22 +00:00
return
}
return app.GetOrCreateRunningCluster(dc.ClusterName, *dc)
2022-05-13 13:30:22 +00:00
}
func runCommandOnDev(ctx context.Context, language, command string) common.CommandResults {
2022-05-14 17:55:00 +00:00
cluster, err := Current.DevelopmentCluster(ctx)
exec := Current.Client().CommandExecutor(ctx)
2022-05-13 13:30:22 +00:00
if err != nil {
return common.CommandResults{
ResultType: "error",
2022-05-14 17:55:00 +00:00
Summary: err.Error(),
2022-05-13 13:30:22 +00:00
}
}
return exec.Execute(cluster.ClusterID, language, command)
}
func RunPythonOnDev(ctx context.Context, command string) common.CommandResults {
return runCommandOnDev(ctx, "python", command)
2022-05-14 17:55:00 +00:00
}
*/
// TODO: Add safe access to p.project and p.project.DevCluster that throws errors if
// the fields are not defined properly
func (p *project) GetDevelopmentClusterId(ctx context.Context) (clusterId string, err error) {
clusterId = p.config.DevCluster.ClusterId
clusterName := p.config.DevCluster.ClusterName
if clusterId != "" {
return
} else if clusterName != "" {
// Add workspaces client on init
return getClusterIdFromClusterName(ctx, p.wsc, clusterName)
} else {
// TODO: Add the project config file location used to error message
err = fmt.Errorf("please define either development cluster's cluster_id or cluster_name in your project config")
return
}
}
2022-09-27 16:58:55 +00:00
func runCommandOnDev(ctx context.Context, language, command string) commands.Results {
clusterId, err := Get(ctx).GetDevelopmentClusterId(ctx)
if err != nil {
2022-09-27 16:58:55 +00:00
return commands.Results{
ResultType: "error",
Summary: err.Error(),
}
}
2022-09-27 16:58:55 +00:00
return Get(ctx).wsc.CommandExecutor.Execute(ctx, clusterId, language, command)
}
2022-09-27 16:58:55 +00:00
func RunPythonOnDev(ctx context.Context, command string) commands.Results {
return runCommandOnDev(ctx, "python", command)
}