diff --git a/cmd/workspace/secrets/overrides.go b/cmd/workspace/secrets/overrides.go index d46284bf6a..5443aca285 100644 --- a/cmd/workspace/secrets/overrides.go +++ b/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 +} diff --git a/internal/acc/debug.go b/internal/acc/debug.go new file mode 100644 index 0000000000..467642e240 --- /dev/null +++ b/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) + } +} diff --git a/internal/acc/helpers.go b/internal/acc/helpers.go new file mode 100644 index 0000000000..aa9902745e --- /dev/null +++ b/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) +} diff --git a/internal/acc/workspace.go b/internal/acc/workspace.go new file mode 100644 index 0000000000..8944e199f9 --- /dev/null +++ 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 +} diff --git a/internal/secrets_test.go b/internal/secrets_test.go index 1fdc48bdde..1e9c86abf2 100644 --- a/internal/secrets_test.go +++ b/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) +} diff --git a/libs/cmdio/io.go b/libs/cmdio/io.go index 8b40294b9f..a60231c098 100644 --- a/libs/cmdio/io.go +++ b/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 {