From f64c44285f2d77cd8a11f3feef6aae3eb59ceefd Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pieter.noordhuis@databricks.com>
Date: Tue, 27 Jun 2023 20:42:29 +0200
Subject: [PATCH] Decode contents by default in workspace export command (#531)

## Changes

Also see #525.

The direct download flag has been removed in newer versions because of
the content type issue.

Instead, we can make the command decode the base64 output when the
output mode is text.

```
$ databricks workspace export /some/path/script.sh
#!/bin/bash
echo "this is a script"
```

## Tests

New integration test.
---
 cmd/workspace/workspace/overrides.go |  3 +++
 internal/workspace_test.go           | 21 +++++++++++++++++++++
 libs/cmdio/render.go                 | 23 +++++++++++++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/cmd/workspace/workspace/overrides.go b/cmd/workspace/workspace/overrides.go
index 0a00ba25..e1b97c59 100644
--- a/cmd/workspace/workspace/overrides.go
+++ b/cmd/workspace/workspace/overrides.go
@@ -8,4 +8,7 @@ func init() {
 	{{header "ID"}}	{{header "Type"}}	{{header "Language"}}	{{header "Path"}}
 	{{range .}}{{green "%d" .ObjectId}}	{{blue "%s" .ObjectType}}	{{cyan "%s" .Language}}	{{.Path|cyan}}
 	{{end}}`)
+
+	// The export command prints the contents of the file to stdout by default.
+	exportCmd.Annotations["template"] = `{{.Content | b64_decode}}`
 }
diff --git a/internal/workspace_test.go b/internal/workspace_test.go
index 3aafffe4..dd26bcf4 100644
--- a/internal/workspace_test.go
+++ b/internal/workspace_test.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"path"
 	"path/filepath"
 	"strings"
 	"testing"
@@ -39,6 +40,26 @@ func TestWorkpaceGetStatusErrorWhenNoArguments(t *testing.T) {
 	assert.Equal(t, "accepts 1 arg(s), received 0", err.Error())
 }
 
+func TestWorkpaceExportPrintsContents(t *testing.T) {
+	t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))
+
+	ctx := context.Background()
+	w := databricks.Must(databricks.NewWorkspaceClient())
+	tmpdir := temporaryWorkspaceDir(t, w)
+	f, err := filer.NewWorkspaceFilesClient(w, tmpdir)
+	require.NoError(t, err)
+
+	// Write file to workspace
+	contents := "#!/usr/bin/bash\necho hello, world\n"
+	err = f.Write(ctx, "file-a", strings.NewReader(contents))
+	require.NoError(t, err)
+
+	// Run export
+	stdout, stderr := RequireSuccessfulRun(t, "workspace", "export", path.Join(tmpdir, "file-a"))
+	assert.Equal(t, contents, stdout.String())
+	assert.Equal(t, "", stderr.String())
+}
+
 func setupWorkspaceImportExportTest(t *testing.T) (context.Context, filer.Filer, string) {
 	t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))
 
diff --git a/libs/cmdio/render.go b/libs/cmdio/render.go
index 7b3da190..d641f61d 100644
--- a/libs/cmdio/render.go
+++ b/libs/cmdio/render.go
@@ -1,6 +1,8 @@
 package cmdio
 
 import (
+	"bytes"
+	"encoding/base64"
 	"encoding/json"
 	"io"
 	"strings"
@@ -93,6 +95,27 @@ func renderTemplate(w io.Writer, tmpl string, v any) error {
 		"pretty_date": func(t time.Time) string {
 			return t.Format("2006-01-02T15:04:05Z")
 		},
+		"b64_encode": func(in string) (string, error) {
+			var out bytes.Buffer
+			enc := base64.NewEncoder(base64.StdEncoding, &out)
+			_, err := enc.Write([]byte(in))
+			if err != nil {
+				return "", err
+			}
+			err = enc.Close()
+			if err != nil {
+				return "", err
+			}
+			return out.String(), nil
+		},
+		"b64_decode": func(in string) (string, error) {
+			dec := base64.NewDecoder(base64.StdEncoding, strings.NewReader(in))
+			out, err := io.ReadAll(dec)
+			if err != nil {
+				return "", err
+			}
+			return string(out), nil
+		},
 	}).Parse(tmpl)
 	if err != nil {
 		return err