From d4fef205e276b931713ae3c07a131d34f8940685 Mon Sep 17 00:00:00 2001 From: Chen Shou Date: Fri, 18 Sep 2026 23:43:46 +0000 Subject: [PATCH] Infer Docker registry host from workspace profile --- .nextchanges/cli/auth-docker-host.md | 1 + acceptance/cmd/auth/docker/help/output.txt | 52 ++++ acceptance/cmd/auth/docker/help/script | 2 + cmd/auth/docker/docker.go | 1 + cmd/auth/docker/docker_configure.go | 95 +++--- cmd/auth/docker/docker_configure_test.go | 289 +++++++++++++++++-- cmd/auth/docker/docker_host.go | 117 ++++++++ cmd/auth/docker/docker_host_test.go | 193 +++++++++++++ cmd/auth/docker/docker_profile.go | 34 +++ libs/dockercredentials/docker_config.go | 34 ++- libs/dockercredentials/docker_config_test.go | 32 ++ libs/dockercredentials/registry.go | 25 +- 12 files changed, 799 insertions(+), 76 deletions(-) create mode 100644 .nextchanges/cli/auth-docker-host.md create mode 100644 cmd/auth/docker/docker_host.go create mode 100644 cmd/auth/docker/docker_host_test.go diff --git a/.nextchanges/cli/auth-docker-host.md b/.nextchanges/cli/auth-docker-host.md new file mode 100644 index 00000000000..7fbb5cce01e --- /dev/null +++ b/.nextchanges/cli/auth-docker-host.md @@ -0,0 +1 @@ +* Deprecate `--region` in `databricks auth docker configure` ahead of its removal in the next release, infer the Artifact Registry region when it is omitted, and add `databricks auth docker host --profile ` to show the profile's registry host and credential-helper status. ([#6781](https://github.com/databricks/cli/pull/6781)) diff --git a/acceptance/cmd/auth/docker/help/output.txt b/acceptance/cmd/auth/docker/help/output.txt index 535b2d1f6e9..780e3e2f538 100644 --- a/acceptance/cmd/auth/docker/help/output.txt +++ b/acceptance/cmd/auth/docker/help/output.txt @@ -42,6 +42,7 @@ Usage: Available Commands: configure (Experimental) Configure Docker authentication for Databricks Artifact Registry + host (Experimental) Show the Databricks Artifact Registry host for a profile token (Experimental) Generate a Docker credential Flags: @@ -57,3 +58,54 @@ Global Flags: --workspace-id string Databricks Workspace ID Use "databricks auth docker [command] --help" for more information about a command. + +>>> [CLI] auth docker configure --help +(Experimental) Configure Docker authentication for Databricks Artifact Registry. + +This command installs docker-credential-databricks and configures Docker to use +it for the selected workspace's Artifact Registry host. If the selected profile +does not already include a workspace_id, the command resolves and saves it so +the Docker helper can map the registry host back to the profile. The registry +region is inferred from the workspace's metastore. Select the workspace with +[PROFILE] or --profile; --host, --account-id, and --workspace-id are not +supported. The deprecated --region flag is retained for compatibility; omit it +because it will be fully removed in the next release. + +Usage: + databricks auth docker configure [PROFILE] [flags] + +Flags: + -h, --help help for configure + --region string Artifact Registry region; we recommend omitting this flag because the region is inferred automatically (DEPRECATED: --region will be fully removed in the next release) + +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 + +>>> [CLI] auth docker host --help +(Experimental) Show the Databricks Artifact Registry host for a profile. + +The --profile flag is required. + +Usage: + databricks auth docker host [flags] + +Examples: + databricks auth docker host --profile DEFAULT + +Flags: + -h, --help help for host + +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 diff --git a/acceptance/cmd/auth/docker/help/script b/acceptance/cmd/auth/docker/help/script index fe76e441c59..1d935e47230 100644 --- a/acceptance/cmd/auth/docker/help/script +++ b/acceptance/cmd/auth/docker/help/script @@ -1,2 +1,4 @@ trace "$CLI" auth --help trace "$CLI" auth docker --help +trace "$CLI" auth docker configure --help +trace "$CLI" auth docker host --help diff --git a/cmd/auth/docker/docker.go b/cmd/auth/docker/docker.go index 5368d8692a9..fa2d3145371 100644 --- a/cmd/auth/docker/docker.go +++ b/cmd/auth/docker/docker.go @@ -17,6 +17,7 @@ func New(load TokenLoader) *cobra.Command { } cmd.AddCommand(newDockerTokenCommand(load)) cmd.AddCommand(newDockerConfigureCommand()) + cmd.AddCommand(newDockerHostCommand()) return cmd } diff --git a/cmd/auth/docker/docker_configure.go b/cmd/auth/docker/docker_configure.go index 2b01a5341f9..dcd8a31cea5 100644 --- a/cmd/auth/docker/docker_configure.go +++ b/cmd/auth/docker/docker_configure.go @@ -4,7 +4,6 @@ import ( "context" "errors" "fmt" - "os" "path/filepath" "strings" @@ -20,24 +19,14 @@ import ( ) type configureDockerDeps struct { - profiler profile.Profiler - newWorkspaceClient func(*databricks.Config) (*databricks.WorkspaceClient, error) - resolveWorkspaceID func(context.Context, *databricks.WorkspaceClient) (string, error) - executable func() (string, error) - registryHost func(string, string, string) (string, error) + dockerProfileDeps installShim func(string) (dockercredentials.ShimInstallResult, error) setCredentialHelper func(string, string) error } func defaultConfigureDockerDeps() configureDockerDeps { return configureDockerDeps{ - profiler: profile.DefaultProfiler, - newWorkspaceClient: func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { - return databricks.NewWorkspaceClient(cfg) - }, - resolveWorkspaceID: authlib.ResolveWorkspaceID, - executable: os.Executable, - registryHost: dockercredentials.RegistryHost, + dockerProfileDeps: defaultDockerProfileDeps(), installShim: dockercredentials.InstallShim, setCredentialHelper: dockercredentials.SetCredentialHelper, } @@ -49,55 +38,83 @@ func newDockerConfigureCommand() *cobra.Command { func newDockerConfigureCommandWithDeps(deps configureDockerDeps) *cobra.Command { cmd := &cobra.Command{ - Use: "configure [PROFILE] --region REGION", + Use: "configure [PROFILE]", Short: "(Experimental) Configure Docker authentication for Databricks Artifact Registry", Long: `(Experimental) Configure Docker authentication for Databricks Artifact Registry. This command installs docker-credential-databricks and configures Docker to use it for the selected workspace's Artifact Registry host. If the selected profile does not already include a workspace_id, the command resolves and saves it so -the Docker helper can map the registry host back to the profile. The required -region must match the workspace home region because it cannot be inferred from -the profile. Select the workspace with [PROFILE] or --profile; --host, ---account-id, and --workspace-id are not supported.`, +the Docker helper can map the registry host back to the profile. The registry +region is inferred from the workspace's metastore. Select the workspace with +[PROFILE] or --profile; --host, --account-id, and --workspace-id are not +supported. The deprecated --region flag is retained for compatibility; omit it +because it will be fully removed in the next release.`, Args: cobra.MaximumNArgs(1), } - var region string - cmd.Flags().StringVar(®ion, "region", "", "Cloud region for the Databricks Artifact Registry host; must match the workspace home region") + var regionFlag string + cmd.Flags().StringVar(®ionFlag, "region", "", "Artifact Registry region; we recommend omitting this flag because the region is inferred automatically") + cmd.Flags().Lookup("region").Deprecated = "--region will be fully removed in the next release" cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() if err := errorOnUnsupportedConfigureDockerFlags(cmd); err != nil { return err } - if region == "" { - return errors.New("--region is required because workspace region cannot be inferred from this profile; it must match the workspace home region") - } profileName, err := configureDockerProfileName(ctx, cmd, args, deps.profiler) if err != nil { return err } - p, err := loadAndValidateConfigureDockerProfile(ctx, profileName, deps.profiler) + p, err := loadAndValidateDockerProfile(ctx, profileName, deps.profiler) if err != nil { return err } + if err := deps.validateWorkspaceHost(p.Host); err != nil { + return err + } + regionProvided := cmd.Flags().Changed("region") + region := strings.TrimSpace(regionFlag) + if regionProvided { + if err := dockercredentials.ValidateRegion(region); err != nil { + return err + } + } executable, err := deps.executable() if err != nil { return fmt.Errorf("locate databricks executable: %w", err) } - workspaceID, err := resolveConfigureDockerWorkspaceID(ctx, p, executable, deps) - if err != nil { - return err + needsWorkspaceClient := !regionProvided || p.WorkspaceID == "" || p.WorkspaceID == authlib.WorkspaceIDNone + var w *databricks.WorkspaceClient + if needsWorkspaceClient { + w, err = newDockerWorkspaceClient(ctx, p, executable, deps.dockerProfileDeps) + if err != nil { + return err + } } - registryHost, err := deps.registryHost(workspaceID, region, p.Host) + workspaceID, err := resolveDockerWorkspaceID(ctx, p, w, deps.dockerProfileDeps) if err != nil { return err } if err := ensureConfigureDockerUniqueProfile(ctx, deps.profiler, p, workspaceID); err != nil { return err } + if !regionProvided { + w.Config.WorkspaceID = workspaceID + region, err = deps.resolveWorkspaceRegion(ctx, w) + if err != nil { + return fmt.Errorf("resolve workspace region for profile %q: %w", p.Name, err) + } + region = strings.TrimSpace(region) + if region == "" { + return fmt.Errorf("resolve workspace region for profile %q: metastore summary did not include a region", p.Name) + } + } + registryHost, err := deps.registryHost(workspaceID, region, p.Host) + if err != nil { + return err + } if p.WorkspaceID == "" || p.WorkspaceID == authlib.WorkspaceIDNone { if err := persistConfigureDockerWorkspaceID(ctx, p, workspaceID); err != nil { return fmt.Errorf("save workspace ID to profile %q: %w", p.Name, err) @@ -108,7 +125,7 @@ the profile. Select the workspace with [PROFILE] or --profile; --host, if err != nil { return fmt.Errorf("install Docker credential helper: %w", err) } - dockerConfigPath, err := configureDockerConfigPath(ctx) + dockerConfigPath, err := dockerConfigPath(ctx) if err != nil { return err } @@ -180,7 +197,7 @@ func configureDockerProfileName(ctx context.Context, cmd *cobra.Command, args [] }) } -func loadAndValidateConfigureDockerProfile(ctx context.Context, profileName string, profiler profile.Profiler) (profile.Profile, error) { +func loadAndValidateDockerProfile(ctx context.Context, profileName string, profiler profile.Profiler) (profile.Profile, error) { profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName)) if err != nil { return profile.Profile{}, err @@ -194,11 +211,7 @@ func loadAndValidateConfigureDockerProfile(ctx context.Context, profileName stri return profiles[0], nil } -func resolveConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, executable string, deps configureDockerDeps) (string, error) { - if p.WorkspaceID != "" && p.WorkspaceID != authlib.WorkspaceIDNone { - return p.WorkspaceID, nil - } - +func newDockerWorkspaceClient(ctx context.Context, p profile.Profile, executable string, deps dockerProfileDeps) (*databricks.WorkspaceClient, error) { cfg := &databricks.Config{ Profile: p.Name, Host: p.Host, @@ -210,8 +223,16 @@ func resolveConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, e } w, err := deps.newWorkspaceClient(cfg) if err != nil { - return "", fmt.Errorf("load workspace profile %q: %w. Run databricks auth login --host and retry with that profile", p.Name, err) + return nil, fmt.Errorf("load workspace profile %q: %w. Run databricks auth login --host and retry with that profile", p.Name, err) } + return w, nil +} + +func resolveDockerWorkspaceID(ctx context.Context, p profile.Profile, w *databricks.WorkspaceClient, deps dockerProfileDeps) (string, error) { + if p.WorkspaceID != "" && p.WorkspaceID != authlib.WorkspaceIDNone { + return p.WorkspaceID, nil + } + // The selected profile may contain the CLI-only "none" sentinel, which the SDK would send as a routing header. w.Config.WorkspaceID = "" workspaceID, err := deps.resolveWorkspaceID(ctx, w) @@ -255,7 +276,7 @@ func persistConfigureDockerWorkspaceID(ctx context.Context, p profile.Profile, w }) } -func configureDockerConfigPath(ctx context.Context) (string, error) { +func dockerConfigPath(ctx context.Context) (string, error) { if dockerConfig := env.Get(ctx, "DOCKER_CONFIG"); dockerConfig != "" { return filepath.Join(dockerConfig, "config.json"), nil } diff --git a/cmd/auth/docker/docker_configure_test.go b/cmd/auth/docker/docker_configure_test.go index 403522e6886..e936a140f1e 100644 --- a/cmd/auth/docker/docker_configure_test.go +++ b/cmd/auth/docker/docker_configure_test.go @@ -1,6 +1,7 @@ package docker import ( + "bytes" "context" "encoding/json" "errors" @@ -15,6 +16,7 @@ import ( "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/dockercredentials" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" @@ -69,6 +71,16 @@ func configureDockerRegistryHostStub(t *testing.T, wantWorkspaceID, wantRegion, } } +func configureDockerRegionStub(region string) func(context.Context, *databricks.WorkspaceClient) (string, error) { + return func(context.Context, *databricks.WorkspaceClient) (string, error) { + return region, nil + } +} + +func allowConfigureDockerWorkspaceHost(string) error { + return nil +} + func writeConfigureDockerExecutable(t *testing.T, dir string) string { t.Helper() name := "databricks" @@ -101,14 +113,27 @@ func TestConfigureDockerCommandWritesDockerConfigAndShim(t *testing.T) { t.Setenv("PATH", binDir) registryHost := "123456789.container.us-west-2.staging.cloud.databricks.test" + server := testserver.New(t) + server.Handle("GET", "/api/2.1/unity-catalog/metastore_summary", func(req testserver.Request) any { + assert.Equal(t, "123456789", req.Headers.Get(authlib.WorkspaceIDHeader)) + return map[string]any{"region": "us-west-2"} + }) deps := defaultConfigureDockerDeps() + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost databricksPath := writeConfigureDockerExecutable(t, binDir) deps.executable = func() (string, error) { return databricksPath, nil } + deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + cfg.Host = server.URL + cfg.Token = "test-token" + cfg.AuthType = "pat" + cfg.Profile = "" + return databricks.NewWorkspaceClient(cfg) + } deps.registryHost = configureDockerRegistryHostStub(t, "123456789", "us-west-2", workspaceHost, registryHost) - cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "--profile", "DEFAULT", "--region", "us-west-2") + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "--profile", "DEFAULT") require.NoError(t, cmd.Execute()) helpers := readCredentialHelpers(t, filepath.Join(dockerDir, "config.json")) @@ -124,30 +149,219 @@ func TestConfigureDockerCommandWritesDockerConfigAndShim(t *testing.T) { assert.Contains(t, stderr.String(), filepath.ToSlash(filepath.Join(dockerDir, "config.json"))) } -func TestConfigureDockerCommandDocumentsRegionRequirement(t *testing.T) { - cmd := newDockerConfigureCommandWithDeps(defaultConfigureDockerDeps()) +func TestConfigureDockerCommandUsesExplicitRegionWithoutWorkspaceClient(t *testing.T) { + t.Parallel() + + ctx := env.Set(cmdio.MockDiscard(t.Context()), "DOCKER_CONFIG", t.TempDir()) + workspaceHost := "https://workspace.cloud.databricks.test" + registryHost := "123456789.container.us-west-2.cloud.databricks.test" + deps := defaultConfigureDockerDeps() + deps.profiler = profile.InMemoryProfiler{Profiles: profile.Profiles{{ + Name: "DEFAULT", + Host: workspaceHost, + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }}} + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.executable = func() (string, error) { return "/usr/local/bin/databricks", nil } + deps.newWorkspaceClient = func(*databricks.Config) (*databricks.WorkspaceClient, error) { + t.Fatal("newWorkspaceClient should not be called") + return nil, nil + } + deps.resolveWorkspaceRegion = func(context.Context, *databricks.WorkspaceClient) (string, error) { + t.Fatal("resolveWorkspaceRegion should not be called") + return "", nil + } + deps.registryHost = configureDockerRegistryHostStub(t, "123456789", "us-west-2", workspaceHost, registryHost) + deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { + return dockercredentials.ShimInstallResult{Path: "/usr/local/bin/docker-credential-databricks", OnPath: true}, nil + } + deps.setCredentialHelper = func(_, host string) error { + assert.Equal(t, registryHost, host) + return nil + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", "us-west-2") + require.NoError(t, cmd.Execute()) +} + +func TestConfigureDockerCommandUsesExplicitRegionWhileResolvingWorkspaceID(t *testing.T) { + t.Parallel() + + ctx := env.Set(cmdio.MockDiscard(t.Context()), "DOCKER_CONFIG", t.TempDir()) + workspaceHost := "https://workspace.cloud.databricks.test" + stopErr := errors.New("stop after registry host") + deps := defaultConfigureDockerDeps() + deps.profiler = profile.InMemoryProfiler{Profiles: profile.Profiles{{ + Name: "DEFAULT", + Host: workspaceHost, + AuthType: authlib.AuthTypeDatabricksCli, + }}} + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.executable = func() (string, error) { return "/usr/local/bin/databricks", nil } + deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + return &databricks.WorkspaceClient{Config: (*config.Config)(cfg)}, nil + } + deps.resolveWorkspaceID = func(context.Context, *databricks.WorkspaceClient) (string, error) { + return "123456789", nil + } + deps.resolveWorkspaceRegion = func(context.Context, *databricks.WorkspaceClient) (string, error) { + return "", errors.New("metastore summary called") + } + deps.registryHost = func(workspaceID, region, host string) (string, error) { + assert.Equal(t, "123456789", workspaceID) + assert.Equal(t, "us-west-2", region) + assert.Equal(t, workspaceHost, host) + return "", stopErr + } - assert.Equal(t, "configure [PROFILE] --region REGION", cmd.Use) - assert.Contains(t, cmd.Flag("region").Usage, "workspace home region") + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", "us-west-2") + err := cmd.Execute() + assert.ErrorIs(t, err, stopErr) +} + +func TestConfigureDockerCommandRejectsInvalidExplicitRegionBeforeWorkspaceClient(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + region string + want string + }{ + {name: "empty", region: "", want: "region is required"}, + {name: "whitespace", region: " ", want: "region is required"}, + {name: "invalid DNS label", region: "-us-west-2", want: `invalid region "-us-west-2"`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + ctx := env.Set(cmdio.MockDiscard(t.Context()), "DOCKER_CONFIG", t.TempDir()) + workspaceClientErr := errors.New("workspace client created") + deps := defaultConfigureDockerDeps() + deps.profiler = profile.InMemoryProfiler{Profiles: profile.Profiles{{ + Name: "DEFAULT", + Host: "https://workspace.cloud.databricks.test", + AuthType: authlib.AuthTypeDatabricksCli, + }}} + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.executable = func() (string, error) { return "/usr/local/bin/databricks", nil } + deps.newWorkspaceClient = func(*databricks.Config) (*databricks.WorkspaceClient, error) { + return nil, workspaceClientErr + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", tt.region) + err := cmd.Execute() + assert.ErrorContains(t, err, tt.want) + assert.NotErrorIs(t, err, workspaceClientErr) + }) + } } -func TestConfigureDockerCommandRequiresRegion(t *testing.T) { +func TestConfigureDockerCommandReturnsMetastoreSummaryErrorBeforeMutation(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) - dir := t.TempDir() - configFile := filepath.Join(dir, ".databrickscfg") + dockerDir := t.TempDir() + server := testserver.New(t) + server.Handle("GET", "/api/2.1/unity-catalog/metastore_summary", func(testserver.Request) any { + return testserver.Response{ + StatusCode: http.StatusInternalServerError, + Body: map[string]any{ + "error_code": "INTERNAL_ERROR", + "message": "summary failed", + }, + } + }) - writeConfigureDockerProfile(t, ctx, configFile, &config.Config{ - Profile: "DEFAULT", + t.Setenv("DOCKER_CONFIG", dockerDir) + deps := defaultConfigureDockerDeps() + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.profiler = profile.InMemoryProfiler{Profiles: profile.Profiles{{ + Name: "DEFAULT", Host: "https://workspace.cloud.databricks.test", WorkspaceID: "123456789", AuthType: authlib.AuthTypeDatabricksCli, + }}} + deps.executable = func() (string, error) { return "/usr/local/bin/databricks", nil } + deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + cfg.Host = server.URL + cfg.Token = "test-token" + cfg.AuthType = "pat" + cfg.Profile = "" + return databricks.NewWorkspaceClient(cfg) + } + deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { + t.Fatal("installShim should not be called") + return dockercredentials.ShimInstallResult{}, nil + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT") + err := cmd.Execute() + assert.ErrorContains(t, err, `resolve workspace region for profile "DEFAULT": summary failed`) + assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) +} + +func TestConfigureDockerCommandRejectsEmptyMetastoreRegionBeforeMutation(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + dockerDir := t.TempDir() + server := testserver.New(t) + server.Handle("GET", "/api/2.1/unity-catalog/metastore_summary", func(testserver.Request) any { + return map[string]any{"metastore_id": "metastore-id"} }) - t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + t.Setenv("DOCKER_CONFIG", dockerDir) + deps := defaultConfigureDockerDeps() + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.profiler = profile.InMemoryProfiler{Profiles: profile.Profiles{{ + Name: "DEFAULT", + Host: "https://workspace.cloud.databricks.test", + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + }}} + deps.executable = func() (string, error) { return "/usr/local/bin/databricks", nil } + deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + cfg.Host = server.URL + cfg.Token = "test-token" + cfg.AuthType = "pat" + cfg.Profile = "" + return databricks.NewWorkspaceClient(cfg) + } + deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { + t.Fatal("installShim should not be called") + return dockercredentials.ShimInstallResult{}, nil + } + + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT") + err := cmd.Execute() + assert.ErrorContains(t, err, `resolve workspace region for profile "DEFAULT": metastore summary did not include a region`) + assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) +} + +func TestConfigureDockerCommandDocumentsRegionDeprecation(t *testing.T) { + t.Parallel() + + cmd := newDockerConfigureCommandWithDeps(defaultConfigureDockerDeps()) + regionFlag := cmd.Flag("region") + + assert.Equal(t, "configure [PROFILE]", cmd.Use) + require.NotNil(t, regionFlag) + assert.Contains(t, regionFlag.Usage, "recommend omitting") + assert.Contains(t, regionFlag.Deprecated, "fully removed in the next release") + assert.False(t, regionFlag.Hidden) +} + +func TestConfigureDockerCommandWarnsWhenRegionIsUsed(t *testing.T) { + t.Parallel() + + deps := defaultConfigureDockerDeps() + deps.profiler = profile.InMemoryProfiler{} + cmd := newDockerConfigureCommandWithDeps(deps) + var stderr bytes.Buffer + cmd.Flags().SetOutput(&stderr) + cmd.SetArgs([]string{"missing", "--region", "us-west-2"}) - cmd := newDockerConfigureTestCommand(ctx, "docker", "configure", "DEFAULT") err := cmd.Execute() - assert.ErrorContains(t, err, "--region is required because workspace region cannot be inferred from this profile; it must match the workspace home region") + assert.ErrorContains(t, err, `profile "missing" not found`) + assert.Contains(t, stderr.String(), "will be fully removed in the next release") } func TestConfigureDockerCommandRejectsAccountOnlyProfile(t *testing.T) { @@ -166,7 +380,7 @@ func TestConfigureDockerCommandRejectsAccountOnlyProfile(t *testing.T) { t.Setenv("DATABRICKS_CONFIG_FILE", configFile) t.Setenv("DOCKER_CONFIG", dockerDir) - cmd := newDockerConfigureTestCommand(ctx, "docker", "configure", "account", "--region", "us-west-2") + cmd := newDockerConfigureTestCommand(ctx, "docker", "configure", "account") err := cmd.Execute() assert.ErrorContains(t, err, "databricks auth login --host ") assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) @@ -201,8 +415,13 @@ func TestConfigureDockerCommandPersistsResolvedWorkspaceID(t *testing.T) { } }) testserver.AddDefaultHandlers(server) + server.Handle("GET", "/api/2.1/unity-catalog/metastore_summary", func(req testserver.Request) any { + assert.Equal(t, "999999", req.Headers.Get(authlib.WorkspaceIDHeader)) + return map[string]any{"region": "us-west-2"} + }) deps := defaultConfigureDockerDeps() + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost databricksPath := writeConfigureDockerExecutable(t, filepath.Join(dir, "bin")) deps.executable = func() (string, error) { return databricksPath, nil @@ -219,7 +438,7 @@ func TestConfigureDockerCommandPersistsResolvedWorkspaceID(t *testing.T) { } deps.registryHost = configureDockerRegistryHostStub(t, "999999", "us-west-2", workspaceHost, "999999.container.us-west-2.gcp.databricks.test") - cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "workspace", "--region", "us-west-2") + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "workspace") require.NoError(t, cmd.Execute()) raw, err := os.ReadFile(configFile) @@ -230,7 +449,7 @@ func TestConfigureDockerCommandPersistsResolvedWorkspaceID(t *testing.T) { assert.Equal(t, dockercredentials.HelperName, helpers["999999.container.us-west-2.gcp.databricks.test"]) } -func TestConfigureDockerCommandRejectsUnsupportedWorkspaceHostBeforeProfileAndDockerConfigMutation(t *testing.T) { +func TestConfigureDockerCommandRejectsUnsupportedWorkspaceHostBeforeWorkspaceRequestOrMutation(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) dir := t.TempDir() configFile := filepath.Join(dir, ".databrickscfg") @@ -250,10 +469,8 @@ func TestConfigureDockerCommandRejectsUnsupportedWorkspaceHostBeforeProfileAndDo deps := defaultConfigureDockerDeps() deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { - return &databricks.WorkspaceClient{Config: (*config.Config)(cfg)}, nil - } - deps.resolveWorkspaceID = func(context.Context, *databricks.WorkspaceClient) (string, error) { - return "123456789", nil + t.Fatal("newWorkspaceClient should not be called") + return nil, nil } deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { t.Fatal("installShim should not be called") @@ -264,7 +481,7 @@ func TestConfigureDockerCommandRejectsUnsupportedWorkspaceHostBeforeProfileAndDo return nil } - cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", "us-west-2") + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT") err = cmd.Execute() assert.ErrorContains(t, err, `"workspace.example.test" is not a supported Databricks workspace host`) after, err := os.ReadFile(configFile) @@ -305,7 +522,7 @@ func TestConfigureDockerCommandRejectsUnsupportedAuthProfiles(t *testing.T) { for _, profileName := range []string{"pat", "m2m", "blank-auth"} { t.Run(profileName, func(t *testing.T) { - cmd := newDockerConfigureTestCommand(ctx, "docker", "configure", profileName, "--region", "us-west-2") + cmd := newDockerConfigureTestCommand(ctx, "docker", "configure", profileName) err := cmd.Execute() assert.ErrorContains(t, err, "requires a profile created by databricks auth login") assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) @@ -330,9 +547,9 @@ func TestConfigureDockerCommandRejectsExplicitInheritedFlags(t *testing.T) { t.Setenv("HOME", filepath.Join(dir, "home")) cases := [][]string{ - {"docker", "configure", "DEFAULT", "--region", "us-west-2", "--host", "https://other.cloud.databricks.test"}, - {"docker", "configure", "DEFAULT", "--region", "us-west-2", "--account-id", "abc"}, - {"docker", "configure", "DEFAULT", "--region", "us-west-2", "--workspace-id", "987654321"}, + {"docker", "configure", "DEFAULT", "--host", "https://other.cloud.databricks.test"}, + {"docker", "configure", "DEFAULT", "--account-id", "abc"}, + {"docker", "configure", "DEFAULT", "--workspace-id", "987654321"}, } for _, args := range cases { @@ -364,6 +581,11 @@ func TestConfigureDockerCommandRejectsAmbiguousWorkspaceIDBeforeDockerConfig(t * t.Setenv("HOME", filepath.Join(dir, "home")) deps := defaultConfigureDockerDeps() + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.resolveWorkspaceRegion = func(context.Context, *databricks.WorkspaceClient) (string, error) { + t.Fatal("resolveWorkspaceRegion should not be called") + return "", nil + } deps.registryHost = func(workspaceID, region, _ string) (string, error) { return workspaceID + ".container." + region + ".cloud.databricks.test", nil } @@ -376,7 +598,7 @@ func TestConfigureDockerCommandRejectsAmbiguousWorkspaceIDBeforeDockerConfig(t * return nil } - cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "one", "--region", "us-west-2") + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "one") err := cmd.Execute() assert.ErrorContains(t, err, "multiple Databricks profiles match workspace ID 123456789") assert.ErrorContains(t, err, "Remove duplicate workspace_id entries") @@ -404,6 +626,8 @@ func TestConfigureDockerCommandRejectsSameWorkspaceIDInDifferentEnvironment(t *t t.Setenv("DATABRICKS_CONFIG_FILE", configFile) deps := defaultConfigureDockerDeps() + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.resolveWorkspaceRegion = configureDockerRegionStub("us-west-2") deps.executable = func() (string, error) { return filepath.Join(dir, "databricks"), nil } @@ -423,7 +647,7 @@ func TestConfigureDockerCommandRejectsSameWorkspaceIDInDifferentEnvironment(t *t return nil } - cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "prod", "--region", "us-west-2") + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "prod") err := cmd.Execute() assert.ErrorContains(t, err, "multiple Databricks profiles match workspace ID 123456789") } @@ -467,6 +691,8 @@ func TestConfigureDockerCommandInstallsShimBeforeDockerConfig(t *testing.T) { t.Setenv("HOME", filepath.Join(dir, "home")) deps := defaultConfigureDockerDeps() + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.resolveWorkspaceRegion = configureDockerRegionStub("us-west-2") deps.executable = func() (string, error) { return "/usr/local/bin/databricks", nil } @@ -479,7 +705,7 @@ func TestConfigureDockerCommandInstallsShimBeforeDockerConfig(t *testing.T) { return nil } - cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", "us-west-2") + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT") err := cmd.Execute() assert.ErrorContains(t, err, "install failed") assert.NoFileExists(t, filepath.Join(dockerDir, "config.json")) @@ -492,6 +718,8 @@ func TestConfigureDockerCommandWarnsAboutPATHAndPATHEXT(t *testing.T) { t.Setenv("DOCKER_CONFIG", t.TempDir()) deps := defaultConfigureDockerDeps() + deps.validateWorkspaceHost = allowConfigureDockerWorkspaceHost + deps.resolveWorkspaceRegion = configureDockerRegionStub("us-west-2") deps.profiler = profile.InMemoryProfiler{Profiles: profile.Profiles{ { Name: "DEFAULT", @@ -503,6 +731,9 @@ func TestConfigureDockerCommandWarnsAboutPATHAndPATHEXT(t *testing.T) { deps.executable = func() (string, error) { return "/usr/local/bin/databricks", nil } + deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + return &databricks.WorkspaceClient{Config: (*config.Config)(cfg)}, nil + } deps.registryHost = configureDockerRegistryHostStub(t, "123456789", "us-west-2", workspaceHost, registryHost) deps.installShim = func(string) (dockercredentials.ShimInstallResult, error) { return dockercredentials.ShimInstallResult{ @@ -514,7 +745,7 @@ func TestConfigureDockerCommandWarnsAboutPATHAndPATHEXT(t *testing.T) { return nil } - cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT", "--region", "us-west-2") + cmd := newDockerConfigureTestCommandWithDeps(ctx, deps, "docker", "configure", "DEFAULT") require.NoError(t, cmd.Execute()) assert.Contains(t, stderr.String(), "PATH") assert.Contains(t, stderr.String(), ".CMD is in PATHEXT on Windows") diff --git a/cmd/auth/docker/docker_host.go b/cmd/auth/docker/docker_host.go new file mode 100644 index 00000000000..0ba449e4e23 --- /dev/null +++ b/cmd/auth/docker/docker_host.go @@ -0,0 +1,117 @@ +package docker + +import ( + "errors" + "fmt" + "strings" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/dockercredentials" + "github.com/spf13/cobra" +) + +type dockerHostDeps struct { + dockerProfileDeps + credentialHelperConfigured func(string, string) (bool, error) +} + +type dockerHostOutput struct { + Host string `json:"host"` + Configured bool `json:"configured"` +} + +func defaultDockerHostDeps() dockerHostDeps { + return dockerHostDeps{ + dockerProfileDeps: defaultDockerProfileDeps(), + credentialHelperConfigured: dockercredentials.CredentialHelperConfigured, + } +} + +func newDockerHostCommand() *cobra.Command { + return newDockerHostCommandWithDeps(defaultDockerHostDeps()) +} + +func newDockerHostCommandWithDeps(deps dockerHostDeps) *cobra.Command { + cmd := &cobra.Command{ + Use: "host", + Short: "(Experimental) Show the Databricks Artifact Registry host for a profile", + Long: `(Experimental) Show the Databricks Artifact Registry host for a profile. + +The --profile flag is required.`, + Example: " databricks auth docker host --profile DEFAULT", + Args: cobra.NoArgs, + Annotations: map[string]string{ + "template": "Registry host: {{.Host}}\nCredential helper configured: {{bool .Configured}}\n", + }, + } + cmd.RunE = func(cmd *cobra.Command, _ []string) error { + ctx := cmd.Context() + if err := errorOnUnsupportedDockerHostFlags(cmd); err != nil { + return err + } + + profileName := strings.TrimSpace(cmd.Flag("profile").Value.String()) + if profileName == "" { + return errors.New("--profile is required for auth docker host") + } + + p, err := loadAndValidateDockerProfile(ctx, profileName, deps.profiler) + if err != nil { + return err + } + if err := deps.validateWorkspaceHost(p.Host); err != nil { + return err + } + + executable, err := deps.executable() + if err != nil { + return fmt.Errorf("locate databricks executable: %w", err) + } + w, err := newDockerWorkspaceClient(ctx, p, executable, deps.dockerProfileDeps) + if err != nil { + return err + } + workspaceID, err := resolveDockerWorkspaceID(ctx, p, w, deps.dockerProfileDeps) + if err != nil { + return err + } + + w.Config.WorkspaceID = workspaceID + region, err := deps.resolveWorkspaceRegion(ctx, w) + if err != nil { + return fmt.Errorf("resolve workspace region for profile %q: %w", p.Name, err) + } + region = strings.TrimSpace(region) + if region == "" { + return fmt.Errorf("resolve workspace region for profile %q: metastore summary did not include a region", p.Name) + } + + registryHost, err := deps.registryHost(workspaceID, region, p.Host) + if err != nil { + return err + } + dockerConfigPath, err := dockerConfigPath(ctx) + if err != nil { + return err + } + configured, err := deps.credentialHelperConfigured(dockerConfigPath, registryHost) + if err != nil { + return err + } + + return cmdio.Render(ctx, dockerHostOutput{ + Host: registryHost, + Configured: configured, + }) + } + return cmd +} + +func errorOnUnsupportedDockerHostFlags(cmd *cobra.Command) error { + for _, name := range []string{"host", "account-id", "workspace-id"} { + if cmd.Flag(name).Changed { + return fmt.Errorf("--%s is not supported for auth docker host. Select the workspace with --profile instead", name) + } + } + return nil +} diff --git a/cmd/auth/docker/docker_host_test.go b/cmd/auth/docker/docker_host_test.go new file mode 100644 index 00000000000..2a703f48d50 --- /dev/null +++ b/cmd/auth/docker/docker_host_test.go @@ -0,0 +1,193 @@ +package docker + +import ( + "bytes" + "context" + "errors" + "io" + "os" + "path/filepath" + "testing" + + cmdroot "github.com/databricks/cli/cmd/root" + authlib "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + dockerHostTestRegistryHost = "123456789.container.us-west-2.cloud.databricks.test" + dockerHostTestWorkspaceHost = "https://workspace.cloud.databricks.test" +) + +func dockerHostTestDeps(t *testing.T, p profile.Profile) dockerHostDeps { + t.Helper() + + deps := defaultDockerHostDeps() + deps.profiler = profile.InMemoryProfiler{Profiles: profile.Profiles{p}} + deps.validateWorkspaceHost = func(string) error { return nil } + deps.executable = func() (string, error) { return "/usr/local/bin/databricks", nil } + deps.newWorkspaceClient = func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + return &databricks.WorkspaceClient{Config: (*config.Config)(cfg)}, nil + } + deps.resolveWorkspaceRegion = func(context.Context, *databricks.WorkspaceClient) (string, error) { + return "us-west-2", nil + } + deps.registryHost = func(workspaceID, region, workspaceHost string) (string, error) { + assert.Equal(t, "123456789", workspaceID) + assert.Equal(t, "us-west-2", region) + assert.Equal(t, dockerHostTestWorkspaceHost, workspaceHost) + return dockerHostTestRegistryHost, nil + } + return deps +} + +func dockerHostTestProfile() profile.Profile { + return profile.Profile{ + Name: "workspace", + Host: dockerHostTestWorkspaceHost, + WorkspaceID: "123456789", + AuthType: authlib.AuthTypeDatabricksCli, + } +} + +func TestDockerHostCommandReportsCredentialHelperStatus(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + helper string + configured string + }{ + {name: "configured", helper: "databricks", configured: "YES"}, + {name: "other helper", helper: "desktop", configured: "NO"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + dockerDir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dockerDir, "config.json"), []byte(`{ + "credHelpers": { + "`+dockerHostTestRegistryHost+`": "`+tt.helper+`" + } +}`), 0o600)) + + ctx := env.Set(t.Context(), "DOCKER_CONFIG", dockerDir) + stdout, err := executeDockerHostCommand(ctx, dockerHostTestDeps(t, dockerHostTestProfile()), "--profile", "workspace") + + require.NoError(t, err) + assert.Equal(t, "Registry host: "+dockerHostTestRegistryHost+"\nCredential helper configured: "+tt.configured+"\n", stdout) + }) + } +} + +func TestDockerHostCommandRendersJSON(t *testing.T) { + t.Parallel() + + ctx := env.Set(t.Context(), "DOCKER_CONFIG", t.TempDir()) + stdout, err := executeDockerHostCommand(ctx, dockerHostTestDeps(t, dockerHostTestProfile()), "--profile", "workspace", "--output", "json") + + require.NoError(t, err) + assert.JSONEq(t, `{ + "host": "123456789.container.us-west-2.cloud.databricks.test", + "configured": false +}`, stdout) +} + +func TestDockerHostCommandRequiresProfileFlag(t *testing.T) { + t.Parallel() + + _, err := executeDockerHostCommand(t.Context(), defaultDockerHostDeps()) + assert.ErrorContains(t, err, "--profile is required for auth docker host") +} + +func TestDockerHostCommandRejectsOtherAuthSelectionFlags(t *testing.T) { + t.Parallel() + + for _, name := range []string{"host", "account-id", "workspace-id"} { + t.Run(name, func(t *testing.T) { + t.Parallel() + + _, err := executeDockerHostCommand(t.Context(), defaultDockerHostDeps(), "--profile", "workspace", "--"+name, "value") + assert.ErrorContains(t, err, "--"+name+" is not supported for auth docker host") + }) + } +} + +func TestDockerHostCommandResolvesMissingWorkspaceIDWithoutPersistingIt(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + profileContents := []byte(`[workspace] +host = ` + dockerHostTestWorkspaceHost + ` +auth_type = databricks-cli +workspace_id = none +`) + require.NoError(t, os.WriteFile(configFile, profileContents, 0o600)) + + p := dockerHostTestProfile() + p.WorkspaceID = authlib.WorkspaceIDNone + deps := dockerHostTestDeps(t, p) + deps.resolveWorkspaceID = func(context.Context, *databricks.WorkspaceClient) (string, error) { + return "123456789", nil + } + deps.resolveWorkspaceRegion = func(_ context.Context, w *databricks.WorkspaceClient) (string, error) { + assert.Equal(t, "123456789", w.Config.WorkspaceID) + return "us-west-2", nil + } + + ctx := env.Set(t.Context(), "DATABRICKS_CONFIG_FILE", configFile) + ctx = env.Set(ctx, "DOCKER_CONFIG", filepath.Join(dir, "docker")) + _, err := executeDockerHostCommand(ctx, deps, "--profile", "workspace") + require.NoError(t, err) + + got, err := os.ReadFile(configFile) + require.NoError(t, err) + assert.Equal(t, string(profileContents), string(got)) +} + +func TestDockerHostCommandReturnsMetastoreError(t *testing.T) { + t.Parallel() + + metastoreErr := errors.New("summary failed") + deps := dockerHostTestDeps(t, dockerHostTestProfile()) + deps.resolveWorkspaceRegion = func(context.Context, *databricks.WorkspaceClient) (string, error) { + return "", metastoreErr + } + + _, err := executeDockerHostCommand(t.Context(), deps, "--profile", "workspace") + assert.ErrorContains(t, err, `resolve workspace region for profile "workspace"`) + assert.ErrorIs(t, err, metastoreErr) +} + +func newDockerHostTestRoot(ctx context.Context, deps dockerHostDeps) *cobra.Command { + ctx = cmdctx.GenerateExecId(ctx) + cmd := cmdroot.New(ctx) + authCmd := &cobra.Command{Use: "auth"} + authCmd.PersistentFlags().String("host", "", "Databricks Host") + authCmd.PersistentFlags().String("account-id", "", "Databricks Account ID") + authCmd.PersistentFlags().String("workspace-id", "", "Databricks Workspace ID") + dockerCmd := &cobra.Command{Use: "docker"} + dockerCmd.AddCommand(newDockerHostCommandWithDeps(deps)) + authCmd.AddCommand(dockerCmd) + cmd.AddCommand(authCmd) + return cmd +} + +func executeDockerHostCommand(ctx context.Context, deps dockerHostDeps, args ...string) (string, error) { + cmd := newDockerHostTestRoot(ctx, deps) + var stdout bytes.Buffer + cmd.SetOut(&stdout) + cmd.SetErr(io.Discard) + cmd.SetArgs(append([]string{"auth", "docker", "host"}, args...)) + err := cmd.Execute() + return stdout.String(), err +} diff --git a/cmd/auth/docker/docker_profile.go b/cmd/auth/docker/docker_profile.go index af056110240..72ba2a0cbc6 100644 --- a/cmd/auth/docker/docker_profile.go +++ b/cmd/auth/docker/docker_profile.go @@ -1,13 +1,47 @@ package docker import ( + "context" "fmt" + "os" authlib "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/cli/libs/dockercredentials" + "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" ) +type dockerProfileDeps struct { + profiler profile.Profiler + newWorkspaceClient func(*databricks.Config) (*databricks.WorkspaceClient, error) + resolveWorkspaceID func(context.Context, *databricks.WorkspaceClient) (string, error) + resolveWorkspaceRegion func(context.Context, *databricks.WorkspaceClient) (string, error) + validateWorkspaceHost func(string) error + executable func() (string, error) + registryHost func(string, string, string) (string, error) +} + +func defaultDockerProfileDeps() dockerProfileDeps { + return dockerProfileDeps{ + profiler: profile.DefaultProfiler, + newWorkspaceClient: func(cfg *databricks.Config) (*databricks.WorkspaceClient, error) { + return databricks.NewWorkspaceClient(cfg) + }, + resolveWorkspaceID: authlib.ResolveWorkspaceID, + resolveWorkspaceRegion: func(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + summary, err := w.Metastores.Summary(ctx) + if err != nil { + return "", err + } + return summary.Region, nil + }, + validateWorkspaceHost: dockercredentials.ValidateWorkspaceHost, + executable: os.Executable, + registryHost: dockercredentials.RegistryHost, + } +} + 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) diff --git a/libs/dockercredentials/docker_config.go b/libs/dockercredentials/docker_config.go index 5f95a2cc6f6..8872681b7e1 100644 --- a/libs/dockercredentials/docker_config.go +++ b/libs/dockercredentials/docker_config.go @@ -12,6 +12,20 @@ import ( // HelperName is the suffix Docker uses to resolve docker-credential-databricks. const HelperName = "databricks" +// CredentialHelperConfigured reports whether registryHost uses docker-credential-databricks. +func CredentialHelperConfigured(path, registryHost string) (bool, error) { + config, err := readDockerConfig(path) + if err != nil { + return false, err + } + + helpers, err := credentialHelpers(path, config) + if err != nil { + return false, err + } + return helpers[registryHost] == HelperName, nil +} + // SetCredentialHelper assigns docker-credential-databricks to registryHost without changing other Docker configuration. // See https://docs.docker.com/reference/cli/docker/login/#credential-helpers. func SetCredentialHelper(path, registryHost string) error { @@ -24,11 +38,9 @@ func SetCredentialHelper(path, registryHost string) error { return err } - helpers := map[string]string{} - if raw, ok := config["credHelpers"]; ok { - if err := json.Unmarshal(raw, &helpers); err != nil { - return fmt.Errorf("read Docker config %s: %w", path, err) - } + helpers, err := credentialHelpers(path, config) + if err != nil { + return err } if helpers == nil { helpers = map[string]string{} @@ -48,6 +60,18 @@ func SetCredentialHelper(path, registryHost string) error { return writeDockerConfig(path, config) } +func credentialHelpers(path string, config map[string]json.RawMessage) (map[string]string, error) { + var helpers map[string]string + raw, ok := config["credHelpers"] + if !ok { + return helpers, nil + } + if err := json.Unmarshal(raw, &helpers); err != nil { + return nil, fmt.Errorf("read Docker config %s: %w", path, err) + } + return helpers, nil +} + // resolveDockerConfigPath follows a config symlink so replacement does not remove the link itself. func resolveDockerConfigPath(path string) (string, error) { info, err := os.Lstat(path) diff --git a/libs/dockercredentials/docker_config_test.go b/libs/dockercredentials/docker_config_test.go index 1949339d4cc..a9a9c6c4ee2 100644 --- a/libs/dockercredentials/docker_config_test.go +++ b/libs/dockercredentials/docker_config_test.go @@ -143,3 +143,35 @@ func TestConfigureDockerCredentialHelperRejectsInvalidJSON(t *testing.T) { err := SetCredentialHelper(path, testRegistryHost) assert.ErrorContains(t, err, "read Docker config") } + +func TestCredentialHelperConfigured(t *testing.T) { + t.Parallel() + + path := filepath.Join(t.TempDir(), "config.json") + require.NoError(t, os.WriteFile(path, []byte(`{ + "credHelpers": { + "999.container.us-east-1.cloud.databricks.test": "databricks", + "registry.example.test": "desktop", + "123.container.us-west-2.cloud.databricks.test": "databricks" + } +}`), 0o600)) + + tests := []struct { + name string + host string + want bool + }{ + {name: "configured", host: "123.container.us-west-2.cloud.databricks.test", want: true}, + {name: "other helper", host: "registry.example.test", want: false}, + {name: "absent", host: "456.container.us-west-2.cloud.databricks.test", want: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := CredentialHelperConfigured(path, tt.host) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} diff --git a/libs/dockercredentials/registry.go b/libs/dockercredentials/registry.go index 34d69a2324b..6db7ccfce4d 100644 --- a/libs/dockercredentials/registry.go +++ b/libs/dockercredentials/registry.go @@ -22,6 +22,24 @@ type Registry struct { Host string } +// ValidateWorkspaceHost checks that a workspace host can be used to derive an Artifact Registry host. +func ValidateWorkspaceHost(workspaceHost string) error { + _, err := registryDNSZoneForWorkspaceHost(workspaceHost) + return err +} + +// ValidateRegion checks that a region can be used as an Artifact Registry hostname label. +func ValidateRegion(region string) error { + region = strings.TrimSpace(region) + if region == "" { + return errors.New("region is required") + } + if !isDNSLabel(region) { + return fmt.Errorf("invalid region %q", region) + } + return nil +} + // RegistryHost builds a registry host in the workspace's cloud and environment DNS zone. func RegistryHost(workspaceID, region, workspaceHost string) (string, error) { workspaceID = strings.TrimSpace(workspaceID) @@ -29,14 +47,11 @@ func RegistryHost(workspaceID, region, workspaceHost string) (string, error) { if workspaceID == "" { return "", errors.New("workspace ID is required") } - if region == "" { - return "", errors.New("region is required") - } if !isDNSLabel(workspaceID) { return "", fmt.Errorf("invalid workspace ID %q", workspaceID) } - if !isDNSLabel(region) { - return "", fmt.Errorf("invalid region %q", region) + if err := ValidateRegion(region); err != nil { + return "", err } dnsZone, err := registryDNSZoneForWorkspaceHost(workspaceHost) if err != nil {