Add Docker credential token command

.nextchanges/cli/docker-token.md

@@ -0,0 +1 @@
+* Add experimental `databricks auth docker token` to generate Docker credentials for Databricks Artifact Registry. ([#6194](https://github.com/databricks/cli/pull/6194))

acceptance/cmd/auth/docker/help/output.txt

@@ -0,0 +1,58 @@
+
+>>> [CLI] auth --help
+Authentication related commands. For more information regarding how
+authentication for the Databricks CLI and SDKs work please refer to the documentation
+linked below.
+
+AWS: https://docs.databricks.com/dev-tools/auth/index.html
+Azure: https://learn.microsoft.com/azure/databricks/dev-tools/auth
+GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html
+
+Usage:
+  databricks auth [command]
+
+Available Commands:
+  describe    Describes the credentials and the source of those credentials, being used by the CLI to authenticate
+  docker      (Experimental) Manage Docker authentication for Databricks Artifact Registry
+  login       Log into a Databricks workspace or account
+  logout      Log out of a Databricks profile
+  profiles    Lists profiles from ~/.databrickscfg
+  switch      Set the default profile
+  token       Get authentication token
+
+Flags:
+      --account-id string     Databricks Account ID
+  -h, --help                  help for auth
+      --host string           Databricks Host
+      --workspace-id string   Databricks Workspace ID
+
+Global Flags:
+      --debug            enable debug logging
+  -o, --output type      output type: text or json (default text)
+  -p, --profile string   ~/.databrickscfg profile
+  -t, --target string    bundle target to use (if applicable)
+
+Use "databricks auth [command] --help" for more information about a command.
+
+>>> [CLI] auth docker --help
+(Experimental) Manage Docker authentication for Databricks Artifact Registry
+
+Usage:
+  databricks auth docker [command]
+
+Available Commands:
+  token       (Experimental) Generate a Docker credential
+
+Flags:
+  -h, --help   help for docker
+
+Global Flags:
+      --account-id string     Databricks Account ID
+      --debug                 enable debug logging
+      --host string           Databricks Host
+  -o, --output type           output type: text or json (default text)
+  -p, --profile string        ~/.databrickscfg profile
+  -t, --target string         bundle target to use (if applicable)
+      --workspace-id string   Databricks Workspace ID
+
+Use "databricks auth docker [command] --help" for more information about a command.

acceptance/cmd/auth/docker/help/script

@@ -0,0 +1,2 @@
+trace "$CLI" auth --help
+trace "$CLI" auth docker --help

acceptance/cmd/auth/docker/script.prepare

@@ -0,0 +1,27 @@
+setup_docker_auth() {
+    export DATABRICKS_HOST_ORIG="$DATABRICKS_HOST"
+    sethome "./home"
+
+    unset DATABRICKS_ACCOUNT_ID
+    unset DATABRICKS_AUTH_TYPE
+    unset DATABRICKS_CLIENT_ID
+    unset DATABRICKS_CLIENT_SECRET
+    unset DATABRICKS_CONFIG_FILE
+    unset DATABRICKS_CONFIG_PROFILE
+    unset DATABRICKS_DISCOVERY_URL
+    unset DATABRICKS_HOST
+    unset DATABRICKS_TOKEN
+    unset DATABRICKS_WORKSPACE_ID
+
+    export DATABRICKS_AUTH_STORAGE=plaintext
+    export DATABRICKS_LOG_LEVEL=error
+    export TEST_DAR_REGISTRY_HOST=123456789.container.us-west-2.cloud.databricks.com
+
+    cat > "$HOME/.databrickscfg" <<ENDCFG
+[docker-test]
+host = $DATABRICKS_HOST_ORIG
+auth_type = databricks-cli
+workspace_id = 123456789
+ENDCFG
+
+}

acceptance/cmd/auth/docker/token/output.txt

@@ -0,0 +1,6 @@
+
+>>> [CLI] auth docker token
+{
+  "Secret": "oauth-token",
+  "Username": "oauthtoken"
+}

acceptance/cmd/auth/docker/token/script

@@ -0,0 +1,6 @@
+setup_docker_auth
+
+mkdir -p "$HOME/.databricks"
+cp token-cache.json "$HOME/.databricks/token-cache.json"
+
+printf '%s\n' "$TEST_DAR_REGISTRY_HOST" | trace "$CLI" auth docker token | jq -S .

acceptance/cmd/auth/docker/token/test.toml

@@ -0,0 +1,7 @@
+Ignore = [
+    "home",
+]
+
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+
+Timeout = "15s"

acceptance/cmd/auth/docker/token/token-cache.json

@@ -0,0 +1,11 @@
+{
+  "version": 1,
+  "tokens": {
+    "docker-test": {
+      "access_token": "cached-access-token",
+      "token_type": "Bearer",
+      "refresh_token": "test-refresh-token",
+      "expiry": "2099-01-01T00:00:00Z"
+    }
+  }
+}

cmd/auth/auth.go

@@ -35,6 +35,7 @@ GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html`,
 	cmd.AddCommand(newLogoutCommand())
 	cmd.AddCommand(newProfilesCommand())
 	cmd.AddCommand(newTokenCommand(&authArguments))
+	cmd.AddCommand(newDockerCommand(&authArguments))
 	cmd.AddCommand(newDescribeCommand())
 	cmd.AddCommand(newSwitchCommand())
 	return cmd

cmd/auth/docker.go

@@ -0,0 +1,54 @@
+package auth
+
+import (
+	"time"
+
+	"github.com/databricks/cli/libs/auth"
+	"github.com/databricks/cli/libs/auth/storage"
+	"github.com/databricks/cli/libs/databrickscfg/profile"
+	"github.com/spf13/cobra"
+)
+
+func newDockerCommand(authArguments *auth.AuthArguments) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "docker",
+		Short: "(Experimental) Manage Docker authentication for Databricks Artifact Registry",
+	}
+	cmd.AddCommand(newDockerTokenCommand(authArguments))
+	return cmd
+}
+
+func newDockerTokenCommand(authArguments *auth.AuthArguments) *cobra.Command {
+	return newDockerTokenCommandWithTokenLoader(authArguments, loadToken)
+}
+
+func newDockerTokenCommandWithTokenLoader(authArguments *auth.AuthArguments, load tokenLoader) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "token",
+		Short: "(Experimental) Generate a Docker credential",
+	}
+
+	var tokenTimeout time.Duration
+	cmd.Flags().DurationVar(&tokenTimeout, "timeout", defaultTimeout, "Timeout for acquiring a token.")
+	var noForceRefresh bool
+	cmd.Flags().BoolVar(&noForceRefresh, "no-force-refresh", false, "Use a valid cached token instead of forcing a refresh.")
+
+	cmd.PreRunE = validateDockerTokenRequest
+	cmd.RunE = func(cmd *cobra.Command, _ []string) error {
+		ctx := cmd.Context()
+		tokenStore, mode, err := storage.ResolveStore(ctx, storage.StorageModeUnknown)
+		if err != nil {
+			return err
+		}
+		return runDockerToken(ctx, cmd, loadTokenArgs{
+			authArguments: authArguments,
+			tokenTimeout:  tokenTimeout,
+			// Docker may reuse one credential for a long upload, so maximize its lifetime by refreshing it by default.
+			forceRefresh: !noForceRefresh,
+			profiler:     profile.DefaultProfiler,
+			tokenStore:   tokenStore,
+			mode:         mode,
+		}, load)
+	}
+	return cmd
+}

cmd/auth/docker_profile.go

@@ -0,0 +1,33 @@
+package auth
+
+import (
+	"fmt"
+
+	authlib "github.com/databricks/cli/libs/auth"
+	"github.com/databricks/cli/libs/databrickscfg/profile"
+	"github.com/databricks/databricks-sdk-go/config"
+)
+
+func validateDockerCredentialProfile(p profile.Profile) error {
+	if p.HasClientCredentials {
+		return fmt.Errorf("profile %q uses client credentials. Docker credential helper requires a profile created by databricks auth login", p.Name)
+	}
+	if p.AuthType != authTypeDatabricksCLI {
+		return fmt.Errorf("profile %q uses auth_type %q. Docker credential helper requires a profile created by databricks auth login", p.Name, p.AuthType)
+	}
+	if isDockerCredentialAccountOnlyProfile(p) {
+		return fmt.Errorf("profile %q does not target a workspace. Run databricks auth login --host <workspace-url> and retry with that profile", p.Name)
+	}
+	return nil
+}
+
+func isDockerCredentialAccountOnlyProfile(p profile.Profile) bool {
+	if p.Host == "" {
+		return true
+	}
+	cfg := &config.Config{Host: p.Host, AccountID: p.AccountID, WorkspaceID: p.WorkspaceID}
+	if authlib.IsClassicAccountHost(cfg.CanonicalHostName()) {
+		return true
+	}
+	return p.AccountID != "" && (p.WorkspaceID == "" || p.WorkspaceID == authlib.WorkspaceIDNone)
+}

cmd/auth/docker_profile_test.go

@@ -0,0 +1,78 @@
+package auth
+
+import (
+	"testing"
+
+	"github.com/databricks/cli/libs/databrickscfg/profile"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestValidateDockerCredentialProfile(t *testing.T) {
+	tests := []struct {
+		name      string
+		profile   profile.Profile
+		wantError string
+	}{
+		{
+			name: "workspace",
+			profile: profile.Profile{
+				Name:        "workspace",
+				Host:        "https://workspace.cloud.databricks.test",
+				WorkspaceID: "123456789",
+				AuthType:    authTypeDatabricksCLI,
+			},
+		},
+		{
+			name: "client credentials",
+			profile: profile.Profile{
+				Name:                 "m2m",
+				Host:                 "https://workspace.cloud.databricks.test",
+				WorkspaceID:          "123456789",
+				HasClientCredentials: true,
+			},
+			wantError: "requires a profile created by databricks auth login",
+		},
+		{
+			name: "unsupported auth type",
+			profile: profile.Profile{
+				Name:        "pat",
+				Host:        "https://workspace.cloud.databricks.test",
+				WorkspaceID: "123456789",
+				AuthType:    "pat",
+			},
+			wantError: "requires a profile created by databricks auth login",
+		},
+		{
+			name: "classic account",
+			profile: profile.Profile{
+				Name:        "account",
+				Host:        "https://accounts.cloud.databricks.test",
+				AccountID:   "account-id",
+				WorkspaceID: "123456789",
+				AuthType:    authTypeDatabricksCLI,
+			},
+			wantError: "does not target a workspace",
+		},
+		{
+			name: "unified account",
+			profile: profile.Profile{
+				Name:      "account",
+				Host:      "https://workspace.cloud.databricks.test",
+				AccountID: "account-id",
+				AuthType:  authTypeDatabricksCLI,
+			},
+			wantError: "does not target a workspace",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := validateDockerCredentialProfile(tt.profile)
+			if tt.wantError == "" {
+				assert.NoError(t, err)
+				return
+			}
+			assert.ErrorContains(t, err, tt.wantError)
+		})
+	}
+}