Fix secrets put-secret command

cmd/workspace/secrets/overrides.go

@@ -1,6 +1,11 @@
 package secrets
 
 import (
+	"encoding/base64"
+	"fmt"
+	"io"
+	"os"
+
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/flags"
@@ -40,15 +45,14 @@ var putSecretCmd = &cobra.Command{
   and cannot exceed 128 characters. The maximum allowed secret value size is 128
   KB. The maximum number of secrets in a given scope is 1000.
 
-  The input fields "string_value" or "bytes_value" specify the type of the
-  secret, which will determine the value returned when the secret value is
-  requested. Exactly one must be specified.
+  The arguments "string-value" or "bytes-value" specify the type of the secret,
+  which will determine the value returned when the secret value is requested.
 
-  Throws RESOURCE_DOES_NOT_EXIST if no such secret scope exists. Throws
-  RESOURCE_LIMIT_EXCEEDED if maximum number of secrets in scope is exceeded.
-  Throws INVALID_PARAMETER_VALUE if the key name or value length is invalid.
-  Throws PERMISSION_DENIED if the user does not have permission to make this
-  API call.`,
+  You can specify the secret value in one of three ways:
+  * Specify the value as a string using the --string-value flag.
+  * Input the secret when prompted interactively (single-line secrets).
+  * Pass the secret via standard input (multi-line secrets).
+  `,
 
 	Annotations: map[string]string{},
 	Args: func(cmd *cobra.Command, args []string) error {
@@ -62,6 +66,13 @@ var putSecretCmd = &cobra.Command{
 	RunE: func(cmd *cobra.Command, args []string) (err error) {
 		ctx := cmd.Context()
 		w := root.WorkspaceClient(ctx)
+
+		bytesValueChanged := cmd.Flags().Changed("bytes-value")
+		stringValueChanged := cmd.Flags().Changed("string-value")
+		if bytesValueChanged && stringValueChanged {
+			return fmt.Errorf("cannot specify both --bytes-value and --string-value")
+		}
+
 		if cmd.Flags().Changed("json") {
 			err = putSecretJson.Unmarshal(&putSecretReq)
 			if err != nil {
@@ -71,12 +82,20 @@ var putSecretCmd = &cobra.Command{
 			putSecretReq.Scope = args[0]
 			putSecretReq.Key = args[1]
 
-			value, err := cmdio.Secret(ctx)
-			if err != nil {
-				return err
+			switch {
+			case bytesValueChanged:
+				// Bytes value set; encode as base64.
+				putSecretReq.BytesValue = base64.StdEncoding.EncodeToString([]byte(putSecretReq.BytesValue))
+			case stringValueChanged:
+				// String value set; nothing to do.
+			default:
+				// Neither is specified; read secret value from stdin.
+				bytes, err := promptSecret(cmd)
+				if err != nil {
+					return err
+				}
+				putSecretReq.BytesValue = base64.StdEncoding.EncodeToString(bytes)
 			}
-
-			putSecretReq.StringValue = value
 		}
 
 		err = w.Secrets.PutSecret(ctx, putSecretReq)
@@ -86,3 +105,17 @@ var putSecretCmd = &cobra.Command{
 		return nil
 	},
 }
+
+func promptSecret(cmd *cobra.Command) ([]byte, error) {
+	// If stdin is a TTY, prompt for the secret.
+	if !cmdio.IsInTTY(cmd.Context()) {
+		return io.ReadAll(os.Stdin)
+	}
+
+	value, err := cmdio.Secret(cmd.Context(), "Please enter your secret value")
+	if err != nil {
+		return nil, err
+	}
+
+	return []byte(value), nil
+}

internal/acc/debug.go

@@ -0,0 +1,42 @@
+package acc
+
+import (
+	"encoding/json"
+	"os"
+	"path"
+	"path/filepath"
+	"testing"
+)
+
+// Detects if test is run from "debug test" feature in VS Code.
+func isInDebug() bool {
+	ex, _ := os.Executable()
+	return path.Base(ex) == "__debug_bin"
+}
+
+// Loads debug environment from ~/.databricks/debug-env.json.
+func loadDebugEnvIfRunFromIDE(t *testing.T, key string) {
+	if !isInDebug() {
+		return
+	}
+	home, err := os.UserHomeDir()
+	if err != nil {
+		t.Fatalf("cannot find user home: %s", err)
+	}
+	raw, err := os.ReadFile(filepath.Join(home, ".databricks/debug-env.json"))
+	if err != nil {
+		t.Fatalf("cannot load ~/.databricks/debug-env.json: %s", err)
+	}
+	var conf map[string]map[string]string
+	err = json.Unmarshal(raw, &conf)
+	if err != nil {
+		t.Fatalf("cannot parse ~/.databricks/debug-env.json: %s", err)
+	}
+	vars, ok := conf[key]
+	if !ok {
+		t.Fatalf("~/.databricks/debug-env.json#%s not configured", key)
+	}
+	for k, v := range vars {
+		os.Setenv(k, v)
+	}
+}

internal/acc/helpers.go

@@ -0,0 +1,35 @@
+package acc
+
+import (
+	"fmt"
+	"math/rand"
+	"os"
+	"strings"
+	"testing"
+	"time"
+)
+
+// GetEnvOrSkipTest proceeds with test only with that env variable.
+func GetEnvOrSkipTest(t *testing.T, name string) string {
+	value := os.Getenv(name)
+	if value == "" {
+		t.Skipf("Environment variable %s is missing", name)
+	}
+	return value
+}
+
+const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
+// RandomName gives random name with optional prefix. e.g. qa.RandomName("tf-")
+func RandomName(prefix ...string) string {
+	rand.Seed(time.Now().UnixNano())
+	randLen := 12
+	b := make([]byte, randLen)
+	for i := range b {
+		b[i] = charset[rand.Intn(randLen)]
+	}
+	if len(prefix) > 0 {
+		return fmt.Sprintf("%s%s", strings.Join(prefix, ""), b)
+	}
+	return string(b)
+}

internal/acc/workspace.go

@@ -0,0 +1,68 @@
+package acc
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/stretchr/testify/require"
+)
+
+type WorkspaceT struct {
+	*testing.T
+
+	W *databricks.WorkspaceClient
+
+	ctx context.Context
+
+	exec *compute.CommandExecutorV2
+}
+
+func WorkspaceTest(t *testing.T) (context.Context, *WorkspaceT) {
+	loadDebugEnvIfRunFromIDE(t, "workspace")
+
+	t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))
+
+	w, err := databricks.NewWorkspaceClient()
+	require.NoError(t, err)
+
+	wt := &WorkspaceT{
+		T: t,
+
+		W: w,
+
+		ctx: context.Background(),
+	}
+
+	return wt.ctx, wt
+}
+
+func (t *WorkspaceT) TestClusterID() string {
+	clusterID := GetEnvOrSkipTest(t.T, "TEST_BRICKS_CLUSTER_ID")
+	err := t.W.Clusters.EnsureClusterIsRunning(t.ctx, clusterID)
+	require.NoError(t, err)
+	return clusterID
+}
+
+func (t *WorkspaceT) RunPython(code string) (string, error) {
+	var err error
+
+	// Create command executor only once per test.
+	if t.exec == nil {
+		t.exec, err = t.W.CommandExecution.Start(t.ctx, t.TestClusterID(), compute.LanguagePython)
+		require.NoError(t, err)
+
+		t.Cleanup(func() {
+			err := t.exec.Destroy(t.ctx)
+			require.NoError(t, err)
+		})
+	}
+
+	results, err := t.exec.Execute(t.ctx, code)
+	require.NoError(t, err)
+	require.NotEqual(t, compute.ResultTypeError, results.ResultType, results.Cause)
+	output, ok := results.Data.(string)
+	require.True(t, ok, "unexpected type %T", results.Data)
+	return output, nil
+}

internal/secrets_test.go

@@ -1,12 +1,98 @@
 package internal
 
 import (
+	"context"
+	"encoding/base64"
+	"fmt"
 	"testing"
 
+	"github.com/databricks/cli/internal/acc"
+	"github.com/databricks/databricks-sdk-go/service/workspace"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestSecretsCreateScopeErrWhenNoArguments(t *testing.T) {
 	_, _, err := RequireErrorRun(t, "secrets", "create-scope")
 	assert.Equal(t, "accepts 1 arg(s), received 0", err.Error())
 }
+
+func temporarySecretScope(ctx context.Context, t *acc.WorkspaceT) string {
+	scope := acc.RandomName("cli-acc-")
+	err := t.W.Secrets.CreateScope(ctx, workspace.CreateScope{
+		Scope: scope,
+	})
+	require.NoError(t, err)
+
+	// Delete the scope after the test.
+	t.Cleanup(func() {
+		err := t.W.Secrets.DeleteScopeByScope(ctx, scope)
+		require.NoError(t, err)
+	})
+
+	return scope
+}
+
+func assertSecretStringValue(t *acc.WorkspaceT, scope, key, expected string) {
+	out, err := t.RunPython(fmt.Sprintf(`
+		import base64
+		value = dbutils.secrets.get(scope="%s", key="%s")
+		encoded_value = base64.b64encode(value.encode('utf-8'))
+		print(encoded_value.decode('utf-8'))
+	`, scope, key))
+	require.NoError(t, err)
+
+	decoded, err := base64.StdEncoding.DecodeString(out)
+	require.NoError(t, err)
+	assert.Equal(t, expected, string(decoded))
+}
+
+func assertSecretBytesValue(t *acc.WorkspaceT, scope, key string, expected []byte) {
+	out, err := t.RunPython(fmt.Sprintf(`
+		import base64
+		value = dbutils.secrets.getBytes(scope="%s", key="%s")
+		encoded_value = base64.b64encode(value)
+		print(encoded_value.decode('utf-8'))
+	`, scope, key))
+	require.NoError(t, err)
+
+	decoded, err := base64.StdEncoding.DecodeString(out)
+	require.NoError(t, err)
+	assert.Equal(t, expected, decoded)
+}
+
+func TestSecretsPutSecretStringValue(tt *testing.T) {
+	ctx, t := acc.WorkspaceTest(tt)
+	scope := temporarySecretScope(ctx, t)
+	key := "test-key"
+	value := "test-value\nwith-newlines\n"
+
+	stdout, stderr := RequireSuccessfulRun(t.T, "secrets", "put-secret", scope, key, "--string-value", value)
+	assert.Empty(t, stdout)
+	assert.Empty(t, stderr)
+
+	assertSecretStringValue(t, scope, key, value)
+	assertSecretBytesValue(t, scope, key, []byte(value))
+}
+
+func TestSecretsPutSecretBytesValue(tt *testing.T) {
+	ctx, t := acc.WorkspaceTest(tt)
+
+	if true {
+		// Uncomment below to run this test in isolation.
+		// To be addressed once none of the commands taint global state.
+		t.Skip("skipping because the test above clobbers global state")
+	}
+
+	scope := temporarySecretScope(ctx, t)
+	key := "test-key"
+	value := []byte{0x00, 0x01, 0x02, 0x03}
+
+	stdout, stderr := RequireSuccessfulRun(t.T, "secrets", "put-secret", scope, key, "--bytes-value", string(value))
+	assert.Empty(t, stdout)
+	assert.Empty(t, stderr)
+
+	// Note: this value cannot be represented as Python string,
+	// so we only check equality through the dbutils.secrets.getBytes API.
+	assertSecretBytesValue(t, scope, key, value)
+}

libs/cmdio/io.go

@@ -174,18 +174,19 @@ func Select[V any](ctx context.Context, names map[string]V, label string) (id st
 	return c.Select(stringNames, label)
 }
 
-func (c *cmdIO) Secret() (value string, err error) {
+func (c *cmdIO) Secret(label string) (value string, err error) {
 	prompt := (promptui.Prompt{
-		Label: "Enter your secrets value",
-		Mask:  '*',
+		Label:       label,
+		Mask:        '*',
+		HideEntered: true,
 	})
 
 	return prompt.Run()
 }
 
-func Secret(ctx context.Context) (value string, err error) {
+func Secret(ctx context.Context, label string) (value string, err error) {
 	c := fromContext(ctx)
-	return c.Secret()
+	return c.Secret(label)
 }
 
 type nopWriteCloser struct {