Akshay/lakebox bugbash followups

cmd/lakebox/api.go

@@ -37,6 +37,24 @@ const lakeboxKeysAPIPath = "/api/2.0/lakebox/ssh-keys"
 // for this API."
 const orgIDHeader = "X-Databricks-Org-Id"
 
+// maxNameBytes mirrors the manager-side `Sandbox.name` length cap. The
+// server measures bytes, not characters, so a name made of emoji or other
+// multi-byte UTF-8 hits the limit in fewer visible characters than the user
+// expects. Mirroring the constant client-side lets us fail fast with the
+// observed byte count instead of paying a round-trip for a 400.
+const maxNameBytes = 256
+
+// validateName returns an error when `name` exceeds the wire limit. The
+// error names the observed byte count so the user can recover in one shot
+// (the original server message just said "exceeds 256 bytes" without saying
+// by how much, which is unhelpful when emoji are in play).
+func validateName(name string) error {
+	if n := len(name); n > maxNameBytes {
+		return fmt.Errorf("--name is %d bytes; limit is %d (emoji and most non-ASCII characters count as 2-4 bytes each)", n, maxNameBytes)
+	}
+	return nil
+}
+
 // lakeboxAPI wraps the SDK ApiClient with workspace-id-aware request headers.
 type lakeboxAPI struct {
 	c *client.DatabricksClient

cmd/lakebox/api_test.go

@@ -0,0 +1,31 @@
+package lakebox
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestValidateNameAcceptsAscii(t *testing.T) {
+	require.NoError(t, validateName(""))
+	require.NoError(t, validateName("my-project"))
+	require.NoError(t, validateName(strings.Repeat("a", 256))) // boundary: exactly the limit
+}
+
+func TestValidateNameRejectsOversize(t *testing.T) {
+	err := validateName(strings.Repeat("a", 257))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "257 bytes")
+	assert.Contains(t, err.Error(), "256")
+}
+
+func TestValidateNameCountsBytesNotRunes(t *testing.T) {
+	// 64 panda emoji = 64 × 4 bytes = 256 bytes — at the limit, OK.
+	require.NoError(t, validateName(strings.Repeat("🐼", 64)))
+	// 65 = 260 bytes, rejected.
+	err := validateName(strings.Repeat("🐼", 65))
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "260 bytes")
+}

cmd/lakebox/config.go

@@ -97,6 +97,9 @@ Examples:
 
 			var name *string
 			if cmd.Flags().Changed("name") {
+				if err := validateName(nameFlag); err != nil {
+					return err
+				}
 				n := nameFlag
 				name = &n
 			}

cmd/lakebox/create.go

@@ -28,18 +28,24 @@ Examples:
   databricks lakebox create --json`,
 		PreRunE: root.MustWorkspaceClient,
 		RunE: func(cmd *cobra.Command, args []string) error {
+			if err := validateName(name); err != nil {
+				return err
+			}
+
 			ctx := cmd.Context()
 			w := cmdctx.WorkspaceClient(ctx)
 			api, err := newLakeboxAPI(w)
 			if err != nil {
 				return err
 			}
 
+			wantJSON := jsonOutput(cmd, outputJSON)
+
 			// In JSON mode, suppress the spinner+ok lines so the only thing
 			// on stdout/stderr that a scripted caller has to parse is the
 			// JSON body itself.
 			var result *createResponse
-			if outputJSON {
+			if wantJSON {
 				result, err = api.create(ctx, name)
 				if err != nil {
 					return fmt.Errorf("failed to create lakebox: %w", err)
@@ -72,12 +78,12 @@ Examples:
 			if shouldSetDefault {
 				if err := setDefault(ctx, profile, result.SandboxID); err != nil {
 					warn(ctx, fmt.Sprintf("Could not save default: %v", err))
-				} else if !outputJSON {
+				} else if !wantJSON {
 					field(ctx, cmd.ErrOrStderr(), "default", result.SandboxID)
 				}
 			}
 
-			if outputJSON {
+			if wantJSON {
 				enc := json.NewEncoder(cmd.OutOrStdout())
 				enc.SetIndent("", "  ")
 				return enc.Encode(result)

cmd/lakebox/list.go

@@ -83,7 +83,7 @@ Example:
 			}
 			_ = setSandboxes(ctx, profile, refs)
 
-			if outputJSON {
+			if jsonOutput(cmd, outputJSON) {
 				enc := json.NewEncoder(cmd.OutOrStdout())
 				enc.SetIndent("", "  ")
 				return enc.Encode(entries)

cmd/lakebox/ssh.go

@@ -11,6 +11,7 @@ import (
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/execv"
+	"github.com/databricks/cli/libs/shellquote"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/spf13/cobra"
 )
@@ -186,9 +187,14 @@ Examples:
 				_ = setGatewayHost(ctx, profile, sandboxGatewayHost)
 			}
 
+			// Spinner runs until exec replaces this process (Linux/macOS)
+			// or until the ssh subprocess returns (Windows execv shim).
+			// We deliberately don't print "Connected" — at this point ssh
+			// hasn't actually handshaken yet, so any success affirmation
+			// here would be a lie that gets contradicted by ssh's own
+			// error output on the failure path.
 			s := spin(ctx, "Connecting to "+cmdio.Bold(ctx, lakeboxID)+"…")
 			defer s.Close()
-			s.ok("Connected to " + cmdio.Bold(ctx, lakeboxID))
 			return execSSHDirect(lakeboxID, host, gatewayPort, keyPath, extraArgs)
 		},
 	}
@@ -201,9 +207,37 @@ Examples:
 
 // execSSHDirect replaces the CLI process with ssh (or simulates that on
 // Windows via execv). All options are passed on the command line, so no
-// ~/.ssh/config entry is required. Extra args are appended after the
-// destination for remote commands or ssh flags.
+// ~/.ssh/config entry is required.
 func execSSHDirect(lakeboxID, host, port, keyPath string, extraArgs []string) error {
+	return execv.Execv(execv.Options{
+		Args: buildSSHArgs(lakeboxID, host, port, keyPath, extraArgs),
+		Env:  os.Environ(),
+	})
+}
+
+// buildSSHArgs assembles the argv we hand to the `ssh` binary.
+//
+// `ssh` concatenates remote-command words with spaces and the remote
+// shell re-parses them. That makes the two natural user shapes
+// incompatible by default:
+//
+//   - Single arg that's already a complete shell command:
+//     `lakebox ssh <id> -- 'echo hi | head -3'`
+//     The user expects the remote shell to split and execute this
+//     string. ssh's "concat with spaces" is a no-op here, so we hand
+//     the arg through untouched.
+//
+//   - Multi-arg exec-style invocation:
+//     `lakebox ssh <id> -- bash -c 'echo hi'`
+//     Cobra splits this into `["bash", "-c", "echo hi"]`. ssh's join
+//     produces `bash -c echo hi` on the wire, which bash re-splits into
+//     `-c=echo` and `$0=hi` — bug F22. We fix that by shell-quoting
+//     each arg before append, so the remote sees `bash -c 'echo hi'`.
+//
+// The heuristic: if there's exactly one extra arg, pass it untouched;
+// otherwise quote every arg. shellquote.BashArg leaves safe args alone,
+// so `ls -la /tmp` round-trips unchanged in the multi-arg path.
+func buildSSHArgs(lakeboxID, host, port, keyPath string, extraArgs []string) []string {
 	args := []string{
 		"ssh",
 		"-i", keyPath,
@@ -215,10 +249,11 @@ func execSSHDirect(lakeboxID, host, port, keyPath string, extraArgs []string) er
 		"-o", "LogLevel=ERROR",
 		fmt.Sprintf("%s@%s", lakeboxID, host),
 	}
-	args = append(args, extraArgs...)
-
-	return execv.Execv(execv.Options{
-		Args: args,
-		Env:  os.Environ(),
-	})
+	if len(extraArgs) == 1 {
+		return append(args, extraArgs[0])
+	}
+	for _, a := range extraArgs {
+		args = append(args, shellquote.BashArg(a))
+	}
+	return args
 }

cmd/lakebox/ssh_test.go

@@ -0,0 +1,101 @@
+package lakebox
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBuildSSHArgsBaseFlags(t *testing.T) {
+	got := buildSSHArgs("happy-panda-1234", "gw.example.test", "2222", "/keys/id", nil)
+	// Sanity: the destination is the last arg when no extras.
+	assert.Equal(t, "happy-panda-1234@gw.example.test", got[len(got)-1])
+	// The base flag set is fixed; spot-check a couple.
+	assert.Contains(t, got, "-i")
+	assert.Contains(t, got, "/keys/id")
+	assert.Contains(t, got, "-p")
+	assert.Contains(t, got, "2222")
+}
+
+func TestBuildSSHArgsQuoting(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		extraArgs []string
+		// expected is what we expect at the end of args, after the destination.
+		expected []string
+	}{
+		{
+			name:      "no extras",
+			extraArgs: nil,
+			expected:  nil,
+		},
+		{
+			name:      "single safe word — passed through",
+			extraArgs: []string{"uname"},
+			expected:  []string{"uname"},
+		},
+		{
+			name:      "single string with spaces — passed through so the remote shell can split it",
+			extraArgs: []string{"echo hello-from-string"},
+			expected:  []string{"echo hello-from-string"},
+		},
+		{
+			name:      "single string with a pipe — passed through to the remote shell",
+			extraArgs: []string{"cat /etc/os-release | head -3"},
+			expected:  []string{"cat /etc/os-release | head -3"},
+		},
+		{
+			name:      "multi safe words — no quoting",
+			extraArgs: []string{"ls", "-la", "/tmp"},
+			expected:  []string{"ls", "-la", "/tmp"},
+		},
+		{
+			name: "bash -c '<cmd>' — third arg gets quoted (this is F22)",
+			// The whole point: without the quoting, the remote shell
+			// re-splits "bash -c echo hi" and bash's -c eats just "echo".
+			extraArgs: []string{"bash", "-c", "echo hi"},
+			expected:  []string{"bash", "-c", "'echo hi'"},
+		},
+		{
+			name:      "arg with single quote — escaped, not lost",
+			extraArgs: []string{"echo", "it's fine"},
+			expected:  []string{"echo", `'it'\''s fine'`},
+		},
+		{
+			name:      "glob pattern — quoted so the remote (not the local) shell expands it",
+			extraArgs: []string{"find", ".", "-name", "*.txt"},
+			expected:  []string{"find", ".", "-name", "'*.txt'"},
+		},
+		{
+			name:      "ssh -o key=val style — gets quoted because of '='; harmless given current arg ordering",
+			extraArgs: []string{"-o", "StrictHostKeyChecking=no"},
+			expected:  []string{"-o", "'StrictHostKeyChecking=no'"},
+		},
+		{
+			name:      "empty arg gets explicit empty quotes (not lost)",
+			extraArgs: []string{"sh", "-c", ""},
+			expected:  []string{"sh", "-c", "''"},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			got := buildSSHArgs("id", "host", "2222", "/k", tc.extraArgs)
+			// Trim everything up to and including the destination so we
+			// can assert just on the appended extras.
+			dest := "id@host"
+			cut := -1
+			for i, a := range got {
+				if a == dest {
+					cut = i + 1
+					break
+				}
+			}
+			assert.GreaterOrEqual(t, cut, 0, "destination not found in args")
+			tail := got[cut:]
+			if tc.expected == nil {
+				assert.Empty(t, tail)
+			} else {
+				assert.Equal(t, tc.expected, tail)
+			}
+		})
+	}
+}

cmd/lakebox/sshkey.go

@@ -90,7 +90,7 @@ Examples:
 				return fmt.Errorf("failed to list ssh keys: %w", err)
 			}
 
-			if outputJSON {
+			if jsonOutput(cmd, outputJSON) {
 				enc := json.NewEncoder(cmd.OutOrStdout())
 				enc.SetIndent("", "  ")
 				return enc.Encode(keys)

cmd/lakebox/status.go

@@ -50,7 +50,7 @@ Example:
 			_ = setGatewayHost(ctx, profile, entry.GatewayHost)
 			_ = upsertSandbox(ctx, profile, entry.SandboxID, entry.Name)
 
-			if outputJSON {
+			if jsonOutput(cmd, outputJSON) {
 				enc := json.NewEncoder(cmd.OutOrStdout())
 				enc.SetIndent("", "  ")
 				return enc.Encode(entry)
@@ -63,6 +63,9 @@ Example:
 				field(ctx, out, "name", entry.Name)
 			}
 			field(ctx, out, "status", status(ctx, entry.Status))
+			if entry.GatewayHost != "" {
+				field(ctx, out, "gateway", cmdio.Faint(ctx, entry.GatewayHost))
+			}
 			if entry.FQDN != "" {
 				field(ctx, out, "fqdn", cmdio.Faint(ctx, entry.FQDN))
 			}

cmd/lakebox/ui.go

@@ -7,8 +7,24 @@ import (
 	"strings"
 
 	"github.com/databricks/cli/libs/cmdio"
+	"github.com/spf13/cobra"
 )
 
+// jsonOutput reports whether the user asked for JSON output, via either
+// the lakebox-local `--json` flag or the framework-global `-o json`. The
+// global flag is the standard shape across other `databricks` commands
+// and the validator already rejects bogus values, but until this helper
+// existed lakebox commands silently ignored it and emitted text anyway.
+func jsonOutput(cmd *cobra.Command, jsonFlag bool) bool {
+	if jsonFlag {
+		return true
+	}
+	if f := cmd.Flag("output"); f != nil && f.Value.String() == "json" {
+		return true
+	}
+	return false
+}
+
 // cmdioSpinner is the subset of *cmdio.spinner's method set we need.
 // Defining the interface locally lets us hold the unexported type as a
 // struct field; cmdio's spinner satisfies it structurally.