steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Improve `workspace import` command by allowing references to local files for content (#793) 

## Changes
This PR makes a few really important QOL improvements to the `workspace
import` command.

They are:
1. Adds the `--file` flag, which allows a user to specify a file to read
the content from.
2. Wraps the most common error first time users of this command will run
into with a helpful hint.
3. Minor changes to the command Use string changing `PATH` ->
`TARGET_PATH`


## Tests
Integration tests. The newly added integration tests that check the
--file flag works as expected for both `SOURCE` and `AUTO` format.

Skipped the other formats because the API behaviour for them is
straightforward.
This commit is contained in:
shreyas-goenka 2023-10-05 14:48:59 +02:00 committed by GitHub
parent 0894584132
commit caade735e3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 124 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
cmd/workspace/workspace/overrides.go
View File

@ -1,7 +1,15 @@
package workspace
import (
"encoding/base64"
"errors"
"fmt"
"net/http"
"os"
"strings"
"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"
)
@ -19,7 +27,43 @@ func exportOverride(exportCmd *cobra.Command, exportReq *workspace.ExportRequest
exportCmd.Annotations["template"] = `{{.Content | b64_decode}}`
}
// 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)
}

80
internal/workspace_test.go
View File

@ -2,6 +2,7 @@ package internal
import (
"context"
"encoding/base64"
"errors"
"io"
"net/http"
@ -14,6 +15,7 @@ import (
"github.com/databricks/cli/libs/filer"
"github.com/databricks/databricks-sdk-go"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/service/workspace"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@ -95,6 +97,12 @@ func assertFilerFileContents(t *testing.T, ctx context.Context, f filer.Filer, p
assert.Contains(t, string(b), content)
}
func assertWorkspaceFileType(t *testing.T, ctx context.Context, f filer.Filer, path string, fileType workspace.ObjectType) {
info, err := f.Stat(ctx, path)
require.NoError(t, err)
assert.Equal(t, fileType, info.Sys().(workspace.ObjectInfo).ObjectType)
}
func TestAccExportDir(t *testing.T) {
ctx, f, sourceDir := setupWorkspaceImportExportTest(t)
targetDir := t.TempDir()
@ -239,3 +247,75 @@ func TestAccImportDirWithOverwriteFlag(t *testing.T) {
assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "hello, world")
assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"python\")")
}
func TestAccImportFileUsingContentFormatSource(t *testing.T) {
ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
// Content = `print(1)`. Uploaded as a notebook by default
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "pyScript"),
"--content", base64.StdEncoding.EncodeToString([]byte("print(1)")), "--language=PYTHON")
assertFilerFileContents(t, ctx, workspaceFiler, "pyScript", "print(1)")
assertWorkspaceFileType(t, ctx, workspaceFiler, "pyScript", workspace.ObjectTypeNotebook)
// Import with content = `# Databricks notebook source\nprint(1)`. Uploaded as a notebook with the content just being print(1)
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "pyNb"),
"--content", base64.StdEncoding.EncodeToString([]byte("`# Databricks notebook source\nprint(1)")),
"--language=PYTHON")
assertFilerFileContents(t, ctx, workspaceFiler, "pyNb", "print(1)")
assertWorkspaceFileType(t, ctx, workspaceFiler, "pyNb", workspace.ObjectTypeNotebook)
}
func TestAccImportFileUsingContentFormatAuto(t *testing.T) {
ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
// Content = `# Databricks notebook source\nprint(1)`. Upload as file if path has no extension.
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "py-nb-as-file"),
"--content", base64.StdEncoding.EncodeToString([]byte("`# Databricks notebook source\nprint(1)")), "--format=AUTO")
assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-file", "# Databricks notebook source\nprint(1)")
assertWorkspaceFileType(t, ctx, workspaceFiler, "py-nb-as-file", workspace.ObjectTypeFile)
// Content = `# Databricks notebook source\nprint(1)`. Upload as notebook if path has py extension
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "py-nb-as-notebook.py"),
"--content", base64.StdEncoding.EncodeToString([]byte("`# Databricks notebook source\nprint(1)")), "--format=AUTO")
assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-notebook", "# Databricks notebook source\nprint(1)")
assertWorkspaceFileType(t, ctx, workspaceFiler, "py-nb-as-notebook", workspace.ObjectTypeNotebook)
// Content = `print(1)`. Upload as file if content is not notebook (even if path has .py extension)
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "not-a-notebook.py"), "--content",
base64.StdEncoding.EncodeToString([]byte("print(1)")), "--format=AUTO")
assertFilerFileContents(t, ctx, workspaceFiler, "not-a-notebook.py", "print(1)")
assertWorkspaceFileType(t, ctx, workspaceFiler, "not-a-notebook.py", workspace.ObjectTypeFile)
}
func TestAccImportFileFormatSource(t *testing.T) {
ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "pyNotebook"), "--file", "./testdata/import_dir/pyNotebook.py", "--language=PYTHON")
assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"python\")")
assertWorkspaceFileType(t, ctx, workspaceFiler, "pyNotebook", workspace.ObjectTypeNotebook)
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "scalaNotebook"), "--file", "./testdata/import_dir/scalaNotebook.scala", "--language=SCALA")
assertFilerFileContents(t, ctx, workspaceFiler, "scalaNotebook", "// Databricks notebook source\nprintln(\"scala\")")
assertWorkspaceFileType(t, ctx, workspaceFiler, "scalaNotebook", workspace.ObjectTypeNotebook)
_, _, err := RequireErrorRun(t, "workspace", "import", filepath.Join(targetDir, "scalaNotebook"), "--file", "./testdata/import_dir/scalaNotebook.scala")
assert.ErrorContains(t, err, "The zip file may not be valid or may be an unsupported version. 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")
}
func TestAccImportFileFormatAuto(t *testing.T) {
ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
// Upload as file if path has no extension
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "py-nb-as-file"), "--file", "./testdata/import_dir/pyNotebook.py", "--format=AUTO")
assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-file", "# Databricks notebook source\nprint(\"python\")")
assertWorkspaceFileType(t, ctx, workspaceFiler, "py-nb-as-file", workspace.ObjectTypeFile)
// Upload as notebook if path has extension
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "py-nb-as-notebook.py"), "--file", "./testdata/import_dir/pyNotebook.py", "--format=AUTO")
assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-notebook", "# Databricks notebook source\nprint(\"python\")")
assertWorkspaceFileType(t, ctx, workspaceFiler, "py-nb-as-notebook", workspace.ObjectTypeNotebook)
// Upload as file if content is not notebook (even if path has .py extension)
RequireSuccessfulRun(t, "workspace", "import", filepath.Join(targetDir, "not-a-notebook.py"), "--file", "./testdata/import_dir/file-a", "--format=AUTO")
assertFilerFileContents(t, ctx, workspaceFiler, "not-a-notebook.py", "hello, world\n")
assertWorkspaceFileType(t, ctx, workspaceFiler, "not-a-notebook.py", workspace.ObjectTypeFile)
}