mirror of https://github.com/databricks/cli.git
98 lines
3.2 KiB
Go
98 lines
3.2 KiB
Go
package workspace
|
|
|
|
import (
|
|
"encoding/base64"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"os"
|
|
"strings"
|
|
|
|
"github.com/databricks/cli/cmd/root"
|
|
"github.com/databricks/cli/libs/cmdio"
|
|
"github.com/databricks/databricks-sdk-go/apierr"
|
|
"github.com/databricks/databricks-sdk-go/service/workspace"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
func listOverride(listCmd *cobra.Command, listReq *workspace.ListWorkspaceRequest) {
|
|
listReq.Path = "/"
|
|
listCmd.Annotations["template"] = cmdio.Heredoc(`
|
|
{{header "ID"}} {{header "Type"}} {{header "Language"}} {{header "Path"}}
|
|
{{range .}}{{green "%d" .ObjectId}} {{blue "%s" .ObjectType}} {{cyan "%s" .Language}} {{.Path|cyan}}
|
|
{{end}}`)
|
|
}
|
|
|
|
func exportOverride(exportCmd *cobra.Command, exportReq *workspace.ExportRequest) {
|
|
// The export command prints the contents of the file to stdout by default.
|
|
exportCmd.Annotations["template"] = `{{.Content | b64_decode}}`
|
|
exportCmd.Use = "export SOURCE_PATH"
|
|
|
|
var filePath string
|
|
exportCmd.Flags().StringVar(&filePath, "file", "", `Path on the local file system to save exported file at.`)
|
|
|
|
exportCmd.RunE = func(cmd *cobra.Command, args []string) error {
|
|
ctx := cmd.Context()
|
|
w := root.WorkspaceClient(ctx)
|
|
if len(args) != 1 {
|
|
return fmt.Errorf("expected to have the absolute path of the object or directory")
|
|
}
|
|
exportReq.Path = args[0]
|
|
|
|
response, err := w.Workspace.Export(ctx, *exportReq)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// Render file content to stdout if no file path is specified.
|
|
if filePath == "" {
|
|
return cmdio.Render(ctx, response)
|
|
}
|
|
b, err := base64.StdEncoding.DecodeString(response.Content)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return os.WriteFile(filePath, b, 0755)
|
|
}
|
|
}
|
|
|
|
// Give better errors / hints for common API errors.
|
|
func wrapImportAPIErrors(err error, importReq *workspace.Import) error {
|
|
apiErr := &apierr.APIError{}
|
|
if !errors.As(err, &apiErr) {
|
|
return err
|
|
}
|
|
isFormatSource := importReq.Format == workspace.ImportFormatSource || importReq.Format == ""
|
|
if isFormatSource && apiErr.StatusCode == http.StatusBadRequest &&
|
|
strings.Contains(apiErr.Message, "The zip file may not be valid or may be an unsupported version.") {
|
|
return fmt.Errorf("%w Hint: Objects imported using format=SOURCE are expected to be zip encoded databricks source notebook(s) by default. Please specify a language using the --language flag if you are trying to import a single uncompressed notebook", err)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func importOverride(importCmd *cobra.Command, importReq *workspace.Import) {
|
|
var filePath string
|
|
importCmd.Use = "import TARGET_PATH"
|
|
importCmd.Flags().StringVar(&filePath, "file", "", `Path of local file to import`)
|
|
importCmd.MarkFlagsMutuallyExclusive("content", "file")
|
|
|
|
originalRunE := importCmd.RunE
|
|
importCmd.RunE = func(cmd *cobra.Command, args []string) error {
|
|
if filePath != "" {
|
|
b, err := os.ReadFile(filePath)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
importReq.Content = base64.StdEncoding.EncodeToString(b)
|
|
}
|
|
err := originalRunE(cmd, args)
|
|
return wrapImportAPIErrors(err, importReq)
|
|
}
|
|
|
|
}
|
|
|
|
func init() {
|
|
listOverrides = append(listOverrides, listOverride)
|
|
exportOverrides = append(exportOverrides, exportOverride)
|
|
importOverrides = append(importOverrides, importOverride)
|
|
}
|