databricks-cli/bundle/deployer/locker.go

package deployer

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/databricks/bricks/utilities"
	"github.com/databricks/databricks-sdk-go"
	"github.com/databricks/databricks-sdk-go/client"
	"github.com/databricks/databricks-sdk-go/service/workspace"
	"github.com/google/uuid"
)

// Locker object enables exclusive access to TargetDir's scope for a client. This
// enables multiple clients to deploy to the same scope (ie TargetDir) in an atomic
// manner
//
// Here are some of the details of the locking protocol used here:
//
// 1. Potentially multiple clients race to create a deploy.lock file in
//    TargetDir/.bundle directory with unique ID. The deploy.lock file
//    is a json file containing the State from the locker
//
// 2. Clients read the remote deploy.lock file and if it's ID matches, the client
//    assumes it has the lock on TargetDir. The client is now free to read/write code
//    asserts and deploy databricks assets scoped under TargetDir
//
// 3. To sidestep clients failing to relinquish a lock during a failed deploy attempt
//    we allow clients to forcefully acquire a lock on TargetDir. However forcefully acquired
//    locks come with the following caveats:
//
//       a.  a forcefully acquired lock does not guarentee exclusive access to
//           TargetDir's scope
//       b.  forcefully acquiring a lock(s) on TargetDir can break the assumption
//           of exclusive access that other clients with non forcefully acquired
//           locks might have
type Locker struct {
	// scope of the locker
	TargetDir string
	// Active == true implies exclusive access to TargetDir for the client.
	// This implication break down if locks are forcefully acquired by a user
	Active bool
	// if locker is active, this information about the locker is uploaded onto
	// the workspace so as to let other clients details about the active locker
	State *LockState
}

type LockState struct {
	// unique identifier for the locker
	ID uuid.UUID
	// last timestamp when locker was active
	AcquisitionTime time.Time
	// Only relevant for active lockers
	// IsForced == true implies the lock was acquired forcefully
	IsForced bool
	// creator of this locker
	User string
}

// don't need to hold lock on TargetDir to read locker state
func GetActiveLockState(ctx context.Context, wsc *databricks.WorkspaceClient, path string) (*LockState, error) {
	bytes, err := utilities.GetRawJsonFileContent(ctx, wsc, path)
	if err != nil {
		return nil, err
	}
	remoteLock := LockState{}
	err = json.Unmarshal(bytes, &remoteLock)
	if err != nil {
		return nil, err
	}
	return &remoteLock, nil
}

// asserts whether lock is held by locker. Returns descriptive error with current
// holder details if locker does not hold the lock
func (locker *Locker) assertLockHeld(ctx context.Context, wsc *databricks.WorkspaceClient) error {
	activeLockState, err := GetActiveLockState(ctx, wsc, locker.RemotePath())
	if err != nil && strings.Contains(err.Error(), "File not found.") {
		return fmt.Errorf("no active lock on target dir: %s", err)
	}
	if err != nil {
		return err
	}
	if activeLockState.ID != locker.State.ID && !activeLockState.IsForced {
		return fmt.Errorf("deploy lock acquired by %s at %v. Use --force to override", activeLockState.User, activeLockState.AcquisitionTime)
	}
	if activeLockState.ID != locker.State.ID && activeLockState.IsForced {
		return fmt.Errorf("deploy lock force acquired by %s at %v. Use --force to override", activeLockState.User, activeLockState.AcquisitionTime)
	}
	return nil
}

// idempotent function since overwrite is set to true
func (locker *Locker) PutFile(ctx context.Context, wsc *databricks.WorkspaceClient, pathToFile string, content []byte) error {
	if !locker.Active {
		return fmt.Errorf("failed to put file. deploy lock not held")
	}
	apiClient, err := client.New(wsc.Config)
	if err != nil {
		return err
	}
	apiPath := fmt.Sprintf(
		"/api/2.0/workspace-files/import-file/%s?overwrite=true",
		strings.TrimLeft(pathToFile, "/"))

	err = apiClient.Do(ctx, http.MethodPost, apiPath, bytes.NewReader(content), nil)
	if err != nil {
		// retry after creating parent dirs
		err = wsc.Workspace.MkdirsByPath(ctx, path.Dir(pathToFile))
		if err != nil {
			return fmt.Errorf("could not mkdir to put file: %s", err)
		}
		err = apiClient.Do(ctx, http.MethodPost, apiPath, bytes.NewReader(content), nil)
	}
	return err
}

func (locker *Locker) GetRawJsonFileContent(ctx context.Context, wsc *databricks.WorkspaceClient, path string) ([]byte, error) {
	if !locker.Active {
		return nil, fmt.Errorf("failed to get file. deploy lock not held")
	}
	return utilities.GetRawJsonFileContent(ctx, wsc, path)
}

func (locker *Locker) Lock(ctx context.Context, wsc *databricks.WorkspaceClient, isForced bool) error {
	newLockerState := LockState{
		ID:              locker.State.ID,
		AcquisitionTime: time.Now(),
		IsForced:        isForced,
		User:            locker.State.User,
	}
	bytes, err := json.Marshal(newLockerState)
	if err != nil {
		return err
	}
	err = utilities.WriteFile(ctx, wsc, locker.RemotePath(), bytes, isForced)
	if err != nil && !strings.Contains(err.Error(), fmt.Sprintf("%s already exists", locker.RemotePath())) {
		return err
	}
	err = locker.assertLockHeld(ctx, wsc)
	if err != nil {
		return err
	}

	locker.State = &newLockerState
	locker.Active = true
	return nil
}

func (locker *Locker) Unlock(ctx context.Context, wsc *databricks.WorkspaceClient) error {
	if !locker.Active {
		return fmt.Errorf("unlock called when lock is not held")
	}
	err := locker.assertLockHeld(ctx, wsc)
	if err != nil {
		return fmt.Errorf("unlock called when lock is not held: %s", err)
	}
	err = wsc.Workspace.Delete(ctx,
		workspace.Delete{
			Path:      locker.RemotePath(),
			Recursive: false,
		},
	)
	if err != nil {
		return err
	}
	locker.Active = false
	return nil
}

func (locker *Locker) RemotePath() string {
	return path.Join(locker.TargetDir, ".bundle/deploy.lock")
}

func CreateLocker(user string, targetDir string) *Locker {
	return &Locker{
		TargetDir: targetDir,
		Active:    false,
		State: &LockState{
			ID:   uuid.New(),
			User: user,
		},
	}
}