From caade735e32d95ff941adb2ff95371b5d8d30fa3 Mon Sep 17 00:00:00 2001 From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com> Date: Thu, 5 Oct 2023 14:48:59 +0200 Subject: [PATCH] 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. --- cmd/workspace/workspace/overrides.go | 44 +++++++++++++++ internal/workspace_test.go | 80 ++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/cmd/workspace/workspace/overrides.go b/cmd/workspace/workspace/overrides.go index 9cae5bef..5c0692d5 100644 --- a/cmd/workspace/workspace/overrides.go +++ b/cmd/workspace/workspace/overrides.go @@ -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) } diff --git a/internal/workspace_test.go b/internal/workspace_test.go index 7110d5c9..6513300e 100644 --- a/internal/workspace_test.go +++ b/internal/workspace_test.go @@ -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) +}