air: drop redundant /Me calls from air logs and air get

.nextchanges/air/7613.md

@@ -0,0 +1 @@
+* `air logs`, `air get`, and `air run --watch` no longer make an extra credential-check request before fetching the run, and `air get` reuses the configured workspace ID instead of looking it up. ([#7613](https://github.com/databricks/cli/pull/7613))

cmd/air/get.go

@@ -8,8 +8,10 @@ import (
 	"strconv"
 
 	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/flags"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/service/iam"
@@ -102,6 +104,18 @@ func authError(ctx context.Context, cmd *cobra.Command, err error) error {
 		fmt.Errorf("failed to verify authentication: %w", err))
 }
 
+// credentialError tells a bad credential apart from missing access after a call
+// fails with 401/403. MustWorkspaceClient only attaches credentials (it does not
+// check a PAT server-side), so the Me() probe runs here, only on failure. It
+// returns the probe's error when the credentials themselves fail, or nil.
+func credentialError(ctx context.Context, w *databricks.WorkspaceClient, err error) error {
+	if !errors.Is(err, apierr.ErrUnauthenticated) && !errors.Is(err, apierr.ErrPermissionDenied) {
+		return nil
+	}
+	_, meErr := w.CurrentUser.Me(ctx, iam.MeRequest{})
+	return meErr
+}
+
 // newGetCommand returns the `air get JOB_RUN_ID` command, which shows status,
 // configuration, and timing details for a specific run.
 func newGetCommand() *cobra.Command {
@@ -139,17 +153,11 @@ JOB_RUN_ID is the job run ID returned by "databricks air run".`,
 				fmt.Errorf("invalid JOB_RUN_ID %q: must be a positive integer", args[0]))
 		}
 
-		// Validate authentication against the workspace before fetching or
-		// rendering anything. MustWorkspaceClient's Config.Authenticate only
-		// attaches credentials (e.g. it does not check a PAT server-side), so
-		// without this a bad credential would surface as a confusing failure
-		// mid-render instead of a clear "not authenticated" error here.
-		if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil {
-			return authError(ctx, cmd, err)
-		}
-
 		run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID})
 		if err != nil {
+			if credErr := credentialError(ctx, w, err); credErr != nil {
+				return authError(ctx, cmd, credErr)
+			}
 			// The backend returns this when the run ID is unknown to the user.
 			if errors.Is(err, apierr.ErrResourceDoesNotExist) {
 				return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false,
@@ -159,11 +167,17 @@ JOB_RUN_ID is the job run ID returned by "databricks air run".`,
 				fmt.Errorf("failed to get status for run %d: %w", runID, err))
 		}
 
-		workspaceID, err := w.CurrentWorkspaceID(ctx)
+		// Prefer the configured workspace ID; only look it up when it is unset.
+		id, err := auth.ResolveWorkspaceID(ctx, w)
 		if err != nil {
 			return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true,
 				fmt.Errorf("failed to get workspace id for run %d: %w", runID, err))
 		}
+		workspaceID, err := strconv.ParseInt(id, 10, 64)
+		if err != nil {
+			return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false,
+				fmt.Errorf("configured workspace_id %q is not a numeric workspace ID", id))
+		}
 
 		data := buildGetData(run)
 		data.DashboardURL = dashboardURL(w.Config.Host, runID, workspaceID)

cmd/air/get_test.go

@@ -6,16 +6,20 @@ import (
 	"errors"
 	"net/http"
 	"net/http/httptest"
+	"strings"
+	"sync/atomic"
 	"testing"
 	"text/template"
 
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/flags"
+	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -57,9 +61,9 @@ func TestGetRunInvalidID(t *testing.T) {
 	assert.Contains(t, err.Error(), "invalid JOB_RUN_ID")
 }
 
-// notFoundGetServer serves the auth probe plus a runs/get that reports the run
-// as missing (400 INVALID_PARAMETER_VALUE, which the SDK maps to
-// ErrResourceDoesNotExist for this path).
+// notFoundGetServer serves a runs/get that reports the run as missing (400
+// INVALID_PARAMETER_VALUE, which the SDK maps to ErrResourceDoesNotExist for
+// this path).
 func notFoundGetServer(t *testing.T) *httptest.Server {
 	t.Helper()
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -68,7 +72,7 @@ func notFoundGetServer(t *testing.T) *httptest.Server {
 			_, _ = w.Write([]byte(`{"error_code":"INVALID_PARAMETER_VALUE","message":"Run 5 does not exist."}`))
 			return
 		}
-		// Me() probe and any other config discovery.
+		// SDK config discovery.
 		_, _ = w.Write([]byte(`{"userName":"u@example.com"}`))
 	}))
 	t.Cleanup(srv.Close)
@@ -88,8 +92,9 @@ func TestGetRunNotFound(t *testing.T) {
 
 func TestGetRunAuthFailed(t *testing.T) {
 	m := mocks.NewMockWorkspaceClient(t)
-	// A genuine auth failure (permission denied) is validated before the run is
-	// fetched, so GetRun is never reached and nothing is rendered.
+	// GetRun's 403 is confirmed as a credential failure by the Me() probe, so
+	// nothing is rendered.
+	m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, mock.Anything).Return(nil, apierr.ErrPermissionDenied)
 	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, apierr.ErrPermissionDenied)
 	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient)
 	cmd := withOutput(newGetCommand(), flags.OutputText)
@@ -104,6 +109,7 @@ func TestGetRunAuthTransient(t *testing.T) {
 	m := mocks.NewMockWorkspaceClient(t)
 	// A transient failure at the auth probe must not be misreported as an auth
 	// error; it surfaces as a retryable internal error instead.
+	m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, mock.Anything).Return(nil, apierr.ErrUnauthenticated)
 	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, errors.New("connection reset"))
 	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient)
 	cmd := withOutput(newGetCommand(), flags.OutputText)
@@ -115,6 +121,22 @@ func TestGetRunAuthTransient(t *testing.T) {
 	assert.NotContains(t, err.Error(), "authentication was not successful")
 }
 
+func TestGetRunDeniedWithValidCredentials(t *testing.T) {
+	m := mocks.NewMockWorkspaceClient(t)
+	// A 403 on the run with credentials the Me() probe accepts is missing access
+	// to the run, not a failed login.
+	m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, mock.Anything).Return(nil, apierr.ErrPermissionDenied)
+	m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(&iam.User{UserName: "u@example.com"}, nil)
+	ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient)
+	cmd := withOutput(newGetCommand(), flags.OutputText)
+	cmd.SetContext(ctx)
+
+	err := cmd.RunE(cmd, []string{"5"})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to get status for run 5")
+	assert.NotContains(t, err.Error(), "authentication was not successful")
+}
+
 func TestAuthError(t *testing.T) {
 	ctx := cmdio.MockDiscard(t.Context())
 	cmd := withOutput(newGetCommand(), flags.OutputText)
@@ -156,6 +178,75 @@ func TestGetRunNotFoundJSON(t *testing.T) {
 	assert.Equal(t, jsonError{Code: "NOT_FOUND", Kind: "NOT_FOUND", Message: "run 5 not found: check the run ID and that it is a job run ID"}, got.Error)
 }
 
+// getRunServer serves a terminal runs/get and a /Me that reports workspace
+// 12345 in the X-Databricks-Org-Id header. *meCalls counts /Me requests.
+func getRunServer(t *testing.T, meCalls *atomic.Int32) *httptest.Server {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch {
+		case r.URL.Path == "/api/2.2/jobs/runs/get":
+			_, _ = w.Write([]byte(`{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}}`))
+		case strings.HasSuffix(r.URL.Path, "/scim/v2/Me"):
+			meCalls.Add(1)
+			w.Header().Set("X-Databricks-Org-Id", "12345")
+			_, _ = w.Write([]byte(`{"userName": "u@example.com"}`))
+		default:
+			// SDK config discovery.
+			_, _ = w.Write([]byte(`{}`))
+		}
+	}))
+	t.Cleanup(srv.Close)
+	return srv
+}
+
+func runGetJSON(t *testing.T, w *databricks.WorkspaceClient) (string, error) {
+	t.Helper()
+	var buf bytes.Buffer
+	ctx := cmdctx.SetWorkspaceClient(t.Context(), w)
+	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", ""))
+	cmd := withOutput(newGetCommand(), flags.OutputJSON)
+	cmd.SetContext(ctx)
+	err := cmd.RunE(cmd, []string{"5"})
+	return buf.String(), err
+}
+
+func TestGetRunURLUsesConfiguredWorkspaceID(t *testing.T) {
+	var meCalls atomic.Int32
+	srv := getRunServer(t, &meCalls)
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: srv.URL, Token: "token", WorkspaceID: "67890"})
+	require.NoError(t, err)
+
+	out, err := runGetJSON(t, w)
+	require.NoError(t, err)
+	assert.Contains(t, out, "/jobs/runs/5?o=67890")
+	assert.Zero(t, meCalls.Load(), "a configured workspace ID needs no /Me lookup")
+}
+
+func TestGetRunURLLooksUpUnsetWorkspaceID(t *testing.T) {
+	// The SDK also reads the workspace ID from the environment.
+	t.Setenv("DATABRICKS_WORKSPACE_ID", "")
+	var meCalls atomic.Int32
+	srv := getRunServer(t, &meCalls)
+
+	out, err := runGetJSON(t, newTestWorkspaceClient(t, srv.URL))
+	require.NoError(t, err)
+	assert.Contains(t, out, "/jobs/runs/5?o=12345")
+	assert.Equal(t, int32(1), meCalls.Load())
+}
+
+func TestGetRunRejectsNonNumericWorkspaceID(t *testing.T) {
+	var meCalls atomic.Int32
+	srv := getRunServer(t, &meCalls)
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: srv.URL, Token: "token", WorkspaceID: "not-a-number"})
+	require.NoError(t, err)
+
+	out, err := runGetJSON(t, w)
+	require.ErrorIs(t, err, root.ErrAlreadyPrinted)
+	var got errorEnvelope
+	require.NoError(t, json.Unmarshal([]byte(out), &got))
+	assert.Equal(t, jsonError{Code: "INVALID_ARGS", Kind: "PERMANENT", Message: `configured workspace_id "not-a-number" is not a numeric workspace ID`}, got.Error)
+}
+
 func TestGetTemplateSweep(t *testing.T) {
 	out := renderGet(t, getData{
 		RunID:  "456",

cmd/air/logs.go

@@ -14,7 +14,6 @@ import (
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
-	"github.com/databricks/databricks-sdk-go/service/iam"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/spf13/cobra"
 )
@@ -146,14 +145,11 @@ JOB_RUN_ID is the job run ID returned by "databricks air run".`,
 func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error {
 	w := cmdctx.WorkspaceClient(ctx)
 
-	// Validate credentials server-side before fetching (MustWorkspaceClient only
-	// attaches them), so a bad token fails clearly here.
-	if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil {
-		return authError(ctx, cmd, err)
-	}
-
 	run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: req.runID})
 	if err != nil {
+		if credErr := credentialError(ctx, w, err); credErr != nil {
+			return authError(ctx, cmd, credErr)
+		}
 		if errors.Is(err, apierr.ErrResourceDoesNotExist) {
 			return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false,
 				fmt.Errorf("run %d not found: check the run ID and that it is a job run ID", req.runID))

cmd/air/logs_test.go

@@ -215,8 +215,9 @@ func TestLogsSkipsNodeValidationWhenNodeCountCannotBeResolved(t *testing.T) {
 	}
 }
 
-// completedRunLogsServer serves the auth probe, a terminal runs/get, and a
-// single page of Bricklens logs (newest-first, as the tail fetch requests).
+// completedRunLogsServer serves a terminal runs/get and a single page of
+// Bricklens logs (newest-first, as the tail fetch requests). A successful fetch
+// never needs the Me() credential probe, so the server fails the test on one.
 func completedRunLogsServer(t *testing.T) *httptest.Server {
 	t.Helper()
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -235,9 +236,12 @@ func completedRunLogsServer(t *testing.T) *httptest.Server {
 				{"time_unix_nano": 2000000000, "body": "line two", "node_index": 0},
 				{"time_unix_nano": 1000000000, "body": "line one", "node_index": 0}
 			]}`))
+		case strings.HasSuffix(r.URL.Path, "/scim/v2/Me"):
+			t.Errorf("unexpected Me() probe on a successful logs fetch")
+			w.WriteHeader(http.StatusInternalServerError)
 		default:
-			// Me() probe and SDK config discovery.
-			_, _ = w.Write([]byte(`{"userName":"u@example.com"}`))
+			// SDK config discovery.
+			_, _ = w.Write([]byte(`{}`))
 		}
 	}))
 	t.Cleanup(srv.Close)

cmd/air/run_watch_test.go

@@ -20,7 +20,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-// watchServer serves submit, the auth probe, a RUNNING-to-terminal transition,
+// watchServer serves submit, the user lookup, a RUNNING-to-terminal transition,
 // and a page of Bricklens logs.
 // resultState is the terminal result the run reports (e.g. SUCCESS, FAILED).
 func watchServer(t *testing.T, resultState string) *httptest.Server {
@@ -54,7 +54,7 @@ func watchServer(t *testing.T, resultState string) *httptest.Server {
 			}
 			_, _ = fmt.Fprintf(w, `{"log_records":[%s,%s]}`, first, second)
 		default:
-			// Me() probe, workspace-id, SDK config discovery.
+			// Me() user lookup, workspace-id, SDK config discovery.
 			_, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`))
 		}
 	}))