From f6a86cb6fa145b9ce72da9beafc6b31c33aea75d Mon Sep 17 00:00:00 2001 From: monalisa Date: Mon, 25 Sep 2023 02:42:51 +0200 Subject: [PATCH 1/6] Add e2e test for export --- internal/workspace_test.go | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/internal/workspace_test.go b/internal/workspace_test.go index dd26bcf41b..0a142a421d 100644 --- a/internal/workspace_test.go +++ b/internal/workspace_test.go @@ -239,3 +239,25 @@ func TestAccImportDirWithOverwriteFlag(t *testing.T) { assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "hello, world") 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)) +} From a70df6260037b3c952fbec7d782f642aa882e1f1 Mon Sep 17 00:00:00 2001 From: monalisa Date: Mon, 25 Sep 2023 02:43:10 +0200 Subject: [PATCH 2/6] Add jupyter casE --- internal/workspace_test.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/internal/workspace_test.go b/internal/workspace_test.go index 0a142a421d..e9bec1c9a0 100644 --- a/internal/workspace_test.go +++ b/internal/workspace_test.go @@ -260,4 +260,11 @@ func TestAccExport(t *testing.T) { 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") } From 3fff4dbd8488ffa8ab0c0667e04c6b780dbcd061 Mon Sep 17 00:00:00 2001 From: monalisa Date: Tue, 3 Oct 2023 16:22:10 +0200 Subject: [PATCH 3/6] add file flag --- cmd/workspace/workspace/overrides.go | 47 ++++++++++++++++++++++++++++ internal/workspace_test.go | 36 ++++++++++++++++++++- 2 files changed, 82 insertions(+), 1 deletion(-) diff --git a/cmd/workspace/workspace/overrides.go b/cmd/workspace/workspace/overrides.go index 9cae5bef54..74d031c097 100644 --- a/cmd/workspace/workspace/overrides.go +++ b/cmd/workspace/workspace/overrides.go @@ -1,6 +1,11 @@ package workspace import ( + "encoding/base64" + "fmt" + "os" + + "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go/service/workspace" "github.com/spf13/cobra" @@ -17,6 +22,48 @@ 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) == 0 { + promptSpinner := cmdio.Spinner(ctx) + promptSpinner <- "No SOURCE_PATH argument specified. Loading names for Workspace drop-down." + names, err := w.Workspace.ObjectInfoPathToObjectIdMap(ctx, workspace.ListWorkspaceRequest{}) + close(promptSpinner) + if err != nil { + return fmt.Errorf("failed to load names for Workspace drop-down. Please manually specify required arguments. Original error: %w", err) + } + id, err := cmdio.Select(ctx, names, "The absolute path of the object or directory") + if err != nil { + return err + } + args = append(args, id) + } + 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) + } } func init() { diff --git a/internal/workspace_test.go b/internal/workspace_test.go index 71603b0179..5500e9e3a7 100644 --- a/internal/workspace_test.go +++ b/internal/workspace_test.go @@ -61,7 +61,7 @@ func TestWorkpaceExportPrintsContents(t *testing.T) { } func setupWorkspaceImportExportTest(t *testing.T) (context.Context, filer.Filer, string) { - t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV")) + // t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV")) ctx := context.Background() w := databricks.Must(databricks.NewWorkspaceClient()) @@ -268,3 +268,37 @@ func TestAccExport(t *testing.T) { 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":`) +} From 888ea8938227a6cd2040eb8968306f7a1db080ba Mon Sep 17 00:00:00 2001 From: monalisa Date: Tue, 3 Oct 2023 16:23:08 +0200 Subject: [PATCH 4/6] - --- internal/workspace_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/workspace_test.go b/internal/workspace_test.go index 5500e9e3a7..21f599cb34 100644 --- a/internal/workspace_test.go +++ b/internal/workspace_test.go @@ -61,7 +61,7 @@ func TestWorkpaceExportPrintsContents(t *testing.T) { } func setupWorkspaceImportExportTest(t *testing.T) (context.Context, filer.Filer, string) { - // t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV")) + t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV")) ctx := context.Background() w := databricks.Must(databricks.NewWorkspaceClient()) From be4f94d883e918218473a653fca3c3736ece2353 Mon Sep 17 00:00:00 2001 From: monalisa Date: Thu, 5 Oct 2023 14:52:41 +0200 Subject: [PATCH 5/6] remove prompt logic --- cmd/workspace/workspace/overrides.go | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/cmd/workspace/workspace/overrides.go b/cmd/workspace/workspace/overrides.go index 74d031c097..3cc0f2ea48 100644 --- a/cmd/workspace/workspace/overrides.go +++ b/cmd/workspace/workspace/overrides.go @@ -30,21 +30,6 @@ func exportOverride(exportCmd *cobra.Command, exportReq *workspace.ExportRequest exportCmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() w := root.WorkspaceClient(ctx) - - if len(args) == 0 { - promptSpinner := cmdio.Spinner(ctx) - promptSpinner <- "No SOURCE_PATH argument specified. Loading names for Workspace drop-down." - names, err := w.Workspace.ObjectInfoPathToObjectIdMap(ctx, workspace.ListWorkspaceRequest{}) - close(promptSpinner) - if err != nil { - return fmt.Errorf("failed to load names for Workspace drop-down. Please manually specify required arguments. Original error: %w", err) - } - id, err := cmdio.Select(ctx, names, "The absolute path of the object or directory") - if err != nil { - return err - } - args = append(args, id) - } if len(args) != 1 { return fmt.Errorf("expected to have the absolute path of the object or directory") } From d7a54a59fee2784e14fa2de601a38719a31f0c7b Mon Sep 17 00:00:00 2001 From: monalisa Date: Thu, 5 Oct 2023 15:14:41 +0200 Subject: [PATCH 6/6] - --- cmd/workspace/workspace/overrides.go | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/cmd/workspace/workspace/overrides.go b/cmd/workspace/workspace/overrides.go index 6a36ecc324..1cac67419f 100644 --- a/cmd/workspace/workspace/overrides.go +++ b/cmd/workspace/workspace/overrides.go @@ -2,17 +2,13 @@ package workspace import ( "encoding/base64" - "fmt" - "os" - "errors" - + "fmt" "net/http" - + "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"