github.com/databricks/cli@v0.203.0/bundle/artifacts/artifacts.go (about)

     1  package artifacts
     2  
     3  import (
     4  	"context"
     5  	"crypto/sha256"
     6  	"encoding/base64"
     7  	"errors"
     8  	"fmt"
     9  	"os"
    10  	"path"
    11  
    12  	"github.com/databricks/cli/bundle"
    13  	"github.com/databricks/cli/bundle/artifacts/whl"
    14  	"github.com/databricks/cli/bundle/config"
    15  	"github.com/databricks/cli/libs/cmdio"
    16  	"github.com/databricks/databricks-sdk-go/service/workspace"
    17  )
    18  
    19  type mutatorFactory = func(name string) bundle.Mutator
    20  
    21  var buildMutators map[config.ArtifactType]mutatorFactory = map[config.ArtifactType]mutatorFactory{
    22  	config.ArtifactPythonWheel: whl.Build,
    23  }
    24  
    25  var uploadMutators map[config.ArtifactType]mutatorFactory = map[config.ArtifactType]mutatorFactory{}
    26  
    27  func getBuildMutator(t config.ArtifactType, name string) bundle.Mutator {
    28  	mutatorFactory, ok := buildMutators[t]
    29  	if !ok {
    30  		mutatorFactory = BasicBuild
    31  	}
    32  
    33  	return mutatorFactory(name)
    34  }
    35  
    36  func getUploadMutator(t config.ArtifactType, name string) bundle.Mutator {
    37  	mutatorFactory, ok := uploadMutators[t]
    38  	if !ok {
    39  		mutatorFactory = BasicUpload
    40  	}
    41  
    42  	return mutatorFactory(name)
    43  }
    44  
    45  // Basic Build defines a general build mutator which builds artifact based on artifact.BuildCommand
    46  type basicBuild struct {
    47  	name string
    48  }
    49  
    50  func BasicBuild(name string) bundle.Mutator {
    51  	return &basicBuild{name: name}
    52  }
    53  
    54  func (m *basicBuild) Name() string {
    55  	return fmt.Sprintf("artifacts.Build(%s)", m.name)
    56  }
    57  
    58  func (m *basicBuild) Apply(ctx context.Context, b *bundle.Bundle) error {
    59  	artifact, ok := b.Config.Artifacts[m.name]
    60  	if !ok {
    61  		return fmt.Errorf("artifact doesn't exist: %s", m.name)
    62  	}
    63  
    64  	cmdio.LogString(ctx, fmt.Sprintf("artifacts.Build(%s): Building...", m.name))
    65  
    66  	out, err := artifact.Build(ctx)
    67  	if err != nil {
    68  		return fmt.Errorf("artifacts.Build(%s): %w, output: %s", m.name, err, out)
    69  	}
    70  	cmdio.LogString(ctx, fmt.Sprintf("artifacts.Build(%s): Build succeeded", m.name))
    71  
    72  	return nil
    73  }
    74  
    75  // Basic Upload defines a general upload mutator which uploads artifact as a library to workspace
    76  type basicUpload struct {
    77  	name string
    78  }
    79  
    80  func BasicUpload(name string) bundle.Mutator {
    81  	return &basicUpload{name: name}
    82  }
    83  
    84  func (m *basicUpload) Name() string {
    85  	return fmt.Sprintf("artifacts.Build(%s)", m.name)
    86  }
    87  
    88  func (m *basicUpload) Apply(ctx context.Context, b *bundle.Bundle) error {
    89  	artifact, ok := b.Config.Artifacts[m.name]
    90  	if !ok {
    91  		return fmt.Errorf("artifact doesn't exist: %s", m.name)
    92  	}
    93  
    94  	if len(artifact.Files) == 0 {
    95  		return fmt.Errorf("artifact source is not configured: %s", m.name)
    96  	}
    97  
    98  	err := uploadArtifact(ctx, artifact, b)
    99  	if err != nil {
   100  		return fmt.Errorf("artifacts.Upload(%s): %w", m.name, err)
   101  	}
   102  
   103  	return nil
   104  }
   105  
   106  func uploadArtifact(ctx context.Context, a *config.Artifact, b *bundle.Bundle) error {
   107  	for i := range a.Files {
   108  		f := &a.Files[i]
   109  		if f.NeedsUpload() {
   110  			filename := path.Base(f.Source)
   111  			cmdio.LogString(ctx, fmt.Sprintf("artifacts.Upload(%s): Uploading...", filename))
   112  			remotePath, err := uploadArtifactFile(ctx, f.Source, b)
   113  			if err != nil {
   114  				return err
   115  			}
   116  			cmdio.LogString(ctx, fmt.Sprintf("artifacts.Upload(%s): Upload succeeded", filename))
   117  
   118  			f.RemotePath = remotePath
   119  		}
   120  	}
   121  
   122  	a.NormalisePaths()
   123  	return nil
   124  }
   125  
   126  // Function to upload artifact file to Workspace
   127  func uploadArtifactFile(ctx context.Context, file string, b *bundle.Bundle) (string, error) {
   128  	raw, err := os.ReadFile(file)
   129  	if err != nil {
   130  		return "", fmt.Errorf("unable to read %s: %w", file, errors.Unwrap(err))
   131  	}
   132  
   133  	uploadPath, err := getUploadBasePath(b)
   134  	if err != nil {
   135  		return "", err
   136  	}
   137  
   138  	fileHash := sha256.Sum256(raw)
   139  	remotePath := path.Join(uploadPath, fmt.Sprintf("%x", fileHash), path.Base(file))
   140  	// Make sure target directory exists.
   141  	err = b.WorkspaceClient().Workspace.MkdirsByPath(ctx, path.Dir(remotePath))
   142  	if err != nil {
   143  		return "", fmt.Errorf("unable to create directory for %s: %w", remotePath, err)
   144  	}
   145  
   146  	// Import to workspace.
   147  	err = b.WorkspaceClient().Workspace.Import(ctx, workspace.Import{
   148  		Path:      remotePath,
   149  		Overwrite: true,
   150  		Format:    workspace.ImportFormatAuto,
   151  		Content:   base64.StdEncoding.EncodeToString(raw),
   152  	})
   153  	if err != nil {
   154  		return "", fmt.Errorf("unable to import %s: %w", remotePath, err)
   155  	}
   156  
   157  	return remotePath, nil
   158  }
   159  
   160  func getUploadBasePath(b *bundle.Bundle) (string, error) {
   161  	artifactPath := b.Config.Workspace.ArtifactsPath
   162  	if artifactPath == "" {
   163  		return "", fmt.Errorf("remote artifact path not configured")
   164  	}
   165  
   166  	return path.Join(artifactPath, ".internal"), nil
   167  }