Add `--file` flag to workspace export command (#794)

This PR:
1. Adds the `--file` flag to the workspace export command. This allows
you to specify a output file to write to.
2. Adds e2e integration tests for the workspace export command

cmd/workspace/workspace/overrides.go

@@ -8,6 +8,7 @@ import (
 	"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"
@@ -25,6 +26,33 @@ func listOverride(listCmd *cobra.Command, listReq *workspace.ListWorkspaceReques
 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.

internal/workspace_test.go

@@ -248,6 +248,69 @@ func TestAccImportDirWithOverwriteFlag(t *testing.T) {
 	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"python\")")
 }
 
+func TestAccExport(t *testing.T) {
+	ctx, f, sourceDir := setupWorkspaceImportExportTest(t)
+
+	var err error
+
+	// Export vanilla file
+	err = f.Write(ctx, "file-a", strings.NewReader("abc"))
+	require.NoError(t, err)
+	stdout, _ := RequireSuccessfulRun(t, "workspace", "export", filepath.Join(sourceDir, "file-a"))
+	b, err := io.ReadAll(&stdout)
+	require.NoError(t, err)
+	assert.Equal(t, "abc", string(b))
+
+	// Export python notebook
+	err = f.Write(ctx, "pyNotebook.py", strings.NewReader("# Databricks notebook source"))
+	require.NoError(t, err)
+	stdout, _ = RequireSuccessfulRun(t, "workspace", "export", filepath.Join(sourceDir, "pyNotebook"))
+	b, err = io.ReadAll(&stdout)
+	require.NoError(t, err)
+	assert.Equal(t, "# Databricks notebook source\n", string(b))
+
+	// Export python notebook as jupyter
+	stdout, _ = RequireSuccessfulRun(t, "workspace", "export", filepath.Join(sourceDir, "pyNotebook"), "--format", "JUPYTER")
+	b, err = io.ReadAll(&stdout)
+	require.NoError(t, err)
+	assert.Contains(t, string(b), `"cells":`, "jupyter notebooks contain the cells field")
+	assert.Contains(t, string(b), `"metadata":`, "jupyter notebooks contain the metadata field")
+}
+
+func TestAccExportWithFileFlag(t *testing.T) {
+	ctx, f, sourceDir := setupWorkspaceImportExportTest(t)
+	localTmpDir := t.TempDir()
+
+	var err error
+
+	// Export vanilla file
+	err = f.Write(ctx, "file-a", strings.NewReader("abc"))
+	require.NoError(t, err)
+	stdout, _ := RequireSuccessfulRun(t, "workspace", "export", filepath.Join(sourceDir, "file-a"), "--file", filepath.Join(localTmpDir, "file.txt"))
+	b, err := io.ReadAll(&stdout)
+	require.NoError(t, err)
+	// Expect nothing to be printed to stdout
+	assert.Equal(t, "", string(b))
+	assertLocalFileContents(t, filepath.Join(localTmpDir, "file.txt"), "abc")
+
+	// Export python notebook
+	err = f.Write(ctx, "pyNotebook.py", strings.NewReader("# Databricks notebook source"))
+	require.NoError(t, err)
+	stdout, _ = RequireSuccessfulRun(t, "workspace", "export", filepath.Join(sourceDir, "pyNotebook"), "--file", filepath.Join(localTmpDir, "pyNb.py"))
+	b, err = io.ReadAll(&stdout)
+	require.NoError(t, err)
+	assert.Equal(t, "", string(b))
+	assertLocalFileContents(t, filepath.Join(localTmpDir, "pyNb.py"), "# Databricks notebook source\n")
+
+	// Export python notebook as jupyter
+	stdout, _ = RequireSuccessfulRun(t, "workspace", "export", filepath.Join(sourceDir, "pyNotebook"), "--format", "JUPYTER", "--file", filepath.Join(localTmpDir, "jupyterNb.ipynb"))
+	b, err = io.ReadAll(&stdout)
+	require.NoError(t, err)
+	assert.Equal(t, "", string(b))
+	assertLocalFileContents(t, filepath.Join(localTmpDir, "jupyterNb.ipynb"), `"cells":`)
+	assertLocalFileContents(t, filepath.Join(localTmpDir, "jupyterNb.ipynb"), `"metadata":`)
+}
+
 func TestAccImportFileUsingContentFormatSource(t *testing.T) {
 	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)