[air] Apply permissions to direct-run jobs and MLflow experiments

acceptance/experimental/air/run-submit/output.txt

@@ -10,8 +10,55 @@ Tip: use --watch when submitting a run to stream logs to your terminal.
 Stream logs after submission using:
   databricks experimental air logs 555
 
-=== the ai_runtime_task carries the code_source_path
->>> print_requests.py //api/2.2/jobs/runs/submit
+=== the ai_runtime_task and additive permission grants
+>>> print_requests.py //api/2.2/jobs/runs/submit //api/2.0/mlflow/experiments/create //api/2.0/permissions --sort
+{
+  "method": "PATCH",
+  "path": "/api/2.0/permissions/experiments/exp-456",
+  "body": {
+    "access_control_list": [
+      {
+        "permission_level": "CAN_MANAGE",
+        "user_name": "alice@example.com"
+      },
+      {
+        "group_name": "data-team",
+        "permission_level": "CAN_READ"
+      },
+      {
+        "permission_level": "CAN_EDIT",
+        "service_principal_name": "training-sp"
+      }
+    ]
+  }
+}
+{
+  "method": "PATCH",
+  "path": "/api/2.0/permissions/jobs/123",
+  "body": {
+    "access_control_list": [
+      {
+        "permission_level": "CAN_MANAGE",
+        "user_name": "alice@example.com"
+      },
+      {
+        "group_name": "data-team",
+        "permission_level": "CAN_VIEW"
+      },
+      {
+        "permission_level": "CAN_MANAGE_RUN",
+        "service_principal_name": "training-sp"
+      }
+    ]
+  }
+}
+{
+  "method": "POST",
+  "path": "/api/2.0/mlflow/experiments/create",
+  "body": {
+    "name": "/Users/[USERNAME]/submit-smoke"
+  }
+}
 {
   "method": "POST",
   "path": "/api/2.2/jobs/runs/submit",

acceptance/experimental/air/run-submit/run.yaml.tmpl

@@ -3,6 +3,13 @@ command: python train.py
 compute:
   accelerator_type: GPU_1xH100
   num_accelerators: 1
+permissions:
+  - user_name: alice@example.com
+    level: CAN_MANAGE
+  - group_name: data-team
+    level: CAN_VIEW
+  - service_principal_name: training-sp
+    level: CAN_MANAGE_RUN
 code_source:
   type: snapshot
   snapshot:

acceptance/experimental/air/run-submit/script

@@ -12,7 +12,7 @@ sed "s/COMMIT_SHA/$(git rev-parse HEAD)/" run.yaml.tmpl > run.yaml
 title "submit with a git code_source"
 trace $CLI experimental air run -f run.yaml
 
-title "the ai_runtime_task carries the code_source_path"
-trace print_requests.py //api/2.2/jobs/runs/submit
+title "the ai_runtime_task and additive permission grants"
+trace print_requests.py //api/2.2/jobs/runs/submit //api/2.0/mlflow/experiments/create //api/2.0/permissions --sort
 
 rm -fr .git

acceptance/experimental/air/run-submit/test.toml

@@ -21,6 +21,33 @@ Response.Body = '''
 {"run_id": 555}
 '''
 
+[[Server]]
+Pattern = "GET /api/2.0/mlflow/experiments/get-by-name"
+Response.StatusCode = 404
+Response.Body = '''
+{"error_code":"RESOURCE_DOES_NOT_EXIST","message":"experiment does not exist"}
+'''
+
+[[Server]]
+Pattern = "POST /api/2.0/mlflow/experiments/create"
+Response.Body = '''
+{"experiment_id":"exp-456"}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.2/jobs/runs/get"
+Response.Body = '''
+{"job_id":123,"run_id":555}
+'''
+
+[[Server]]
+Pattern = "PATCH /api/2.0/permissions/jobs/123"
+Response.Body = '{}'
+
+[[Server]]
+Pattern = "PATCH /api/2.0/permissions/experiments/exp-456"
+Response.Body = '{}'
+
 # The snapshot tarball is named <dir>_<cachekey[:16]>.tar.gz, where <dir> is the
 # test's temp-dir basename and the cache key derives from the pinned commit SHA.
 # Both are stable given the pinned commit dates in the script, but the temp-dir

experimental/air/cmd/runpermissions.go

@@ -0,0 +1,164 @@
+package aircmd
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"github.com/databricks/cli/libs/log"
+	"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/databricks/databricks-sdk-go/service/ml"
+)
+
+// permissionExperimentName returns the full MLflow experiment path for a run.
+func permissionExperimentName(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig) (string, error) {
+	if cfg.MLflowExperimentDirectory != nil {
+		return strings.TrimRight(*cfg.MLflowExperimentDirectory, "/") + "/" + cfg.ExperimentName, nil
+	}
+
+	email, err := currentUserEmail(ctx, w)
+	if err != nil {
+		return "", err
+	}
+	return "/Users/" + email + "/" + cfg.ExperimentName, nil
+}
+
+// getOrCreateMLflowExperiment resolves the experiment ID, creating it when absent.
+func getOrCreateMLflowExperiment(ctx context.Context, w *databricks.WorkspaceClient, name, artifactLocation string) (string, error) {
+	existing, err := w.Experiments.GetByName(ctx, ml.GetByNameRequest{ExperimentName: name})
+	if err == nil && existing.Experiment != nil && existing.Experiment.ExperimentId != "" {
+		return existing.Experiment.ExperimentId, nil
+	}
+	if err != nil && !errors.Is(err, apierr.ErrNotFound) {
+		return "", fmt.Errorf("failed to get MLflow experiment %q: %w", name, err)
+	}
+
+	created, err := w.Experiments.CreateExperiment(ctx, ml.CreateExperiment{
+		Name:             name,
+		ArtifactLocation: artifactLocation,
+	})
+	if err == nil {
+		return created.ExperimentId, nil
+	}
+	if !errors.Is(err, apierr.ErrAlreadyExists) && !errors.Is(err, apierr.ErrResourceAlreadyExists) {
+		return "", fmt.Errorf("failed to create MLflow experiment %q: %w", name, err)
+	}
+
+	existing, err = w.Experiments.GetByName(ctx, ml.GetByNameRequest{ExperimentName: name})
+	if err != nil {
+		return "", fmt.Errorf("failed to get concurrently created MLflow experiment %q: %w", name, err)
+	}
+	if existing.Experiment == nil || existing.Experiment.ExperimentId == "" {
+		return "", fmt.Errorf("MLflow experiment %q exists but has no experiment ID", name)
+	}
+	return existing.Experiment.ExperimentId, nil
+}
+
+// experimentPermissionLevel maps a Jobs permission level to its MLflow equivalent.
+func experimentPermissionLevel(level string) (iam.PermissionLevel, error) {
+	switch iam.PermissionLevel(level) {
+	case iam.PermissionLevelCanView:
+		return iam.PermissionLevelCanRead, nil
+	case iam.PermissionLevelCanManageRun:
+		return iam.PermissionLevelCanEdit, nil
+	case iam.PermissionLevelCanManage, iam.PermissionLevelIsOwner:
+		return iam.PermissionLevelCanManage, nil
+	default:
+		return "", fmt.Errorf("unsupported AIR permission level %q", level)
+	}
+}
+
+// permissionAccessControl builds an ACL entry for a validated permission grant.
+func permissionAccessControl(p permission, level iam.PermissionLevel) iam.AccessControlRequest {
+	acl := iam.AccessControlRequest{PermissionLevel: level}
+	switch {
+	case p.UserName != nil:
+		acl.UserName = *p.UserName
+	case p.GroupName != nil:
+		acl.GroupName = *p.GroupName
+	case p.ServicePrincipalName != nil:
+		acl.ServicePrincipalName = *p.ServicePrincipalName
+	}
+	return acl
+}
+
+// grantWorkloadPermissions adds the configured ACLs to a job and its experiment.
+func grantWorkloadPermissions(ctx context.Context, w *databricks.WorkspaceClient, jobID, experimentID string, permissions []permission) error {
+	if len(permissions) == 0 {
+		return nil
+	}
+
+	jobACL := make([]iam.AccessControlRequest, 0, len(permissions))
+	experimentACL := make([]iam.AccessControlRequest, 0, len(permissions))
+	for _, p := range permissions {
+		experimentLevel, err := experimentPermissionLevel(p.Level)
+		if err != nil {
+			return err
+		}
+		jobACL = append(jobACL, permissionAccessControl(p, iam.PermissionLevel(p.Level)))
+		experimentACL = append(experimentACL, permissionAccessControl(p, experimentLevel))
+	}
+
+	_, err := w.Permissions.Update(ctx, iam.UpdateObjectPermissions{
+		RequestObjectType: "jobs",
+		RequestObjectId:   jobID,
+		AccessControlList: jobACL,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to grant job permissions: %w", err)
+	}
+
+	_, err = w.Permissions.Update(ctx, iam.UpdateObjectPermissions{
+		RequestObjectType: "experiments",
+		RequestObjectId:   experimentID,
+		AccessControlList: experimentACL,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to grant MLflow experiment permissions: %w", err)
+	}
+	return nil
+}
+
+// preparePermissionExperiment resolves the experiment before the workload is submitted.
+func preparePermissionExperiment(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig) string {
+	if len(cfg.Permissions) == 0 {
+		return ""
+	}
+
+	name, err := permissionExperimentName(ctx, w, cfg)
+	if err != nil {
+		log.Warnf(ctx, "unable to resolve MLflow experiment name; skipping permission grants: %v", err)
+		return ""
+	}
+	artifactLocation := ""
+	if cfg.MLflowArtifactLocation != nil {
+		artifactLocation = *cfg.MLflowArtifactLocation
+	}
+	experimentID, err := getOrCreateMLflowExperiment(ctx, w, name, artifactLocation)
+	if err != nil {
+		log.Warnf(ctx, "unable to get or create MLflow experiment; skipping permission grants: %v", err)
+		return ""
+	}
+	return experimentID
+}
+
+// applySubmittedPermissions resolves the submitted job and adds its configured ACLs.
+func applySubmittedPermissions(ctx context.Context, w *databricks.WorkspaceClient, runID int64, experimentID string, permissions []permission) {
+	if len(permissions) == 0 || experimentID == "" {
+		return
+	}
+
+	run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID})
+	if err == nil {
+		err = grantWorkloadPermissions(ctx, w, strconv.FormatInt(run.JobId, 10), experimentID, permissions)
+	}
+	if err != nil {
+		log.Warnf(ctx, "failed to grant permissions on workload: %v", err)
+		log.Warnf(ctx, "job was created successfully, but permissions could not be granted")
+	}
+}

experimental/air/cmd/runpermissions_test.go

@@ -0,0 +1,113 @@
+package aircmd
+
+import (
+	"encoding/json"
+	"net/http"
+	"testing"
+
+	"github.com/databricks/cli/libs/testserver"
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/iam"
+	"github.com/databricks/databricks-sdk-go/service/ml"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetOrCreateMLflowExperimentCreatesWithArtifactLocation(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	server.Handle("GET", "/api/2.0/mlflow/experiments/get-by-name", func(req testserver.Request) any {
+		return testserver.Response{
+			StatusCode: http.StatusNotFound,
+			Body: map[string]string{
+				"error_code": "RESOURCE_DOES_NOT_EXIST",
+				"message":    "experiment does not exist",
+			},
+		}
+	})
+	server.Handle("POST", "/api/2.0/mlflow/experiments/create", func(req testserver.Request) any {
+		var got ml.CreateExperiment
+		require.NoError(t, json.Unmarshal(req.Body, &got))
+		assert.Equal(t, "/Users/alice@example.com/training", got.Name)
+		assert.Equal(t, "dbfs:/Volumes/main/default/artifacts", got.ArtifactLocation)
+		return ml.CreateExperimentResponse{ExperimentId: "exp-456"}
+	})
+
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
+	require.NoError(t, err)
+	experimentID, err := getOrCreateMLflowExperiment(t.Context(), w, "/Users/alice@example.com/training", "dbfs:/Volumes/main/default/artifacts")
+	require.NoError(t, err)
+	assert.Equal(t, "exp-456", experimentID)
+}
+
+func TestExperimentPermissionLevel(t *testing.T) {
+	tests := []struct {
+		job        string
+		experiment iam.PermissionLevel
+	}{
+		{"CAN_VIEW", iam.PermissionLevelCanRead},
+		{"CAN_MANAGE_RUN", iam.PermissionLevelCanEdit},
+		{"CAN_MANAGE", iam.PermissionLevelCanManage},
+		{"IS_OWNER", iam.PermissionLevelCanManage},
+	}
+	for _, tt := range tests {
+		t.Run(tt.job, func(t *testing.T) {
+			got, err := experimentPermissionLevel(tt.job)
+			require.NoError(t, err)
+			assert.Equal(t, tt.experiment, got)
+		})
+	}
+}
+
+func TestGrantWorkloadPermissionsRejectsUnsupportedLevelBeforeRequests(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
+	require.NoError(t, err)
+	server.RequestCallback = func(req *testserver.Request) {
+		t.Errorf("unexpected permission request: %s %s", req.Method, req.URL.Path)
+	}
+	err = grantWorkloadPermissions(t.Context(), w, "123", "exp-456", []permission{
+		{GroupName: new("data-team"), Level: "CAN_USE"},
+	})
+	require.EqualError(t, err, `unsupported AIR permission level "CAN_USE"`)
+}
+
+func TestGrantWorkloadPermissionsUpdatesJobAndExperiment(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	requests := make(map[string]string)
+	for _, objectPath := range []string{"jobs/123", "experiments/exp-456"} {
+		server.Handle("PATCH", "/api/2.0/permissions/"+objectPath, func(req testserver.Request) any {
+			requests[objectPath] = string(req.Body)
+			return map[string]any{}
+		})
+	}
+
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
+	require.NoError(t, err)
+	err = grantWorkloadPermissions(t.Context(), w, "123", "exp-456", []permission{
+		{UserName: new("alice@example.com"), Level: "CAN_MANAGE"},
+		{GroupName: new("data-team"), Level: "CAN_VIEW"},
+		{ServicePrincipalName: new("training-sp"), Level: "CAN_MANAGE_RUN"},
+	})
+	require.NoError(t, err)
+
+	assert.JSONEq(t, `{
+		"access_control_list": [
+			{"user_name": "alice@example.com", "permission_level": "CAN_MANAGE"},
+			{"group_name": "data-team", "permission_level": "CAN_VIEW"},
+			{"service_principal_name": "training-sp", "permission_level": "CAN_MANAGE_RUN"}
+		]
+	}`, requests["jobs/123"])
+	assert.JSONEq(t, `{
+		"access_control_list": [
+			{"user_name": "alice@example.com", "permission_level": "CAN_MANAGE"},
+			{"group_name": "data-team", "permission_level": "CAN_READ"},
+			{"service_principal_name": "training-sp", "permission_level": "CAN_EDIT"}
+		]
+	}`, requests["experiments/exp-456"])
+}