Improve workspace import command by allowing references to local files for content

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)
 }

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)
+}