Fix secrets put-secret command (#545)

## Changes

Two issues with this command:
* The command line arguments for the secret value were ignored
* If the secret value was piped through stdin, it would still prompt

The second issue prevented users from using multi-line strings because
the prompt reads until end-of-line.

This change adds testing infrastructure for:
* Setting up a workspace focused test (common between many tests)
* Running a snippet of Python through the command execution API

Porting more integration tests to use this infrastructure will be done
in later commits.

## Tests

New integration test passes.

The interactive path cannot be integration tested just yet.
This commit is contained in:
Pieter Noordhuis 2023-07-05 17:30:54 +02:00 committed by GitHub
parent 335475095e
commit db6313e99c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 283 additions and 18 deletions

View File

@ -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
}

42
internal/acc/debug.go Normal file
View File

@ -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)
}
}

35
internal/acc/helpers.go Normal file
View File

@ -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)
}

68
internal/acc/workspace.go Normal file
View File

@ -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
}

View File

@ -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)
}

View File

@ -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 {