diff --git a/cmd/workspace/workspace/overrides.go b/cmd/workspace/workspace/overrides.go index 5c0692d5..1cac6741 100644 --- a/cmd/workspace/workspace/overrides.go +++ b/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. diff --git a/internal/workspace_test.go b/internal/workspace_test.go index 6513300e..39760ec5 100644 --- a/internal/workspace_test.go +++ b/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)