Add SetDefaultHostMetadataResolverFactory

NEXT_CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### New Features and Improvements
 
+ * Add `config.DefaultHostMetadataResolverFactory`: a package-level variable consulted when `Config.HostMetadataResolver` is unset. Lets programs install a shared resolver (e.g. a caching one) once from an `init()` block (typically in a blank-imported package) instead of wiring per-Config. Experimental.
+
 ### Bug Fixes
 
  * Add `X-Databricks-Org-Id` header to `Workspace.Download()` and `Workspace.Upload()` for SPOG host compatibility.

config/config.go

@@ -701,10 +701,17 @@ func (c *Config) resolveHostMetadata(ctx context.Context) {
 		return
 	}
 
+	resolver := c.HostMetadataResolver
+	if resolver == nil {
+		if factory := DefaultHostMetadataResolverFactory; factory != nil {
+			resolver = factory(c)
+		}
+	}
+
 	var meta *HostMetadata
 	var err error
-	if c.HostMetadataResolver != nil {
-		meta, err = c.HostMetadataResolver(ctx, c.CanonicalHostName())
+	if resolver != nil {
+		meta, err = resolver(ctx, c.CanonicalHostName())
 	} else {
 		meta, err = getHostMetadata(ctx, c.CanonicalHostName(), c.refreshClient)
 	}

config/config_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/databricks/databricks-sdk-go/common/environment"
@@ -1169,3 +1170,79 @@ func TestConfig_ResolveHostMetadata_HostTypes(t *testing.T) {
 		})
 	}
 }
+
+// withDefaultHostMetadataResolverFactory installs factory for the duration of
+// the current test, restoring whatever was previously set on cleanup.
+// Capture/set/restore are not atomic — do not use with t.Parallel across
+// multiple tests that touch the package-level default.
+func withDefaultHostMetadataResolverFactory(t *testing.T, factory func(*Config) HostMetadataResolver) {
+	t.Helper()
+	prev := DefaultHostMetadataResolverFactory
+	DefaultHostMetadataResolverFactory = factory
+	t.Cleanup(func() { DefaultHostMetadataResolverFactory = prev })
+}
+
+func TestDefaultHostMetadataResolverFactory_UsedWhenConfigHasNoResolver(t *testing.T) {
+	var factoryCalls atomic.Int32
+	withDefaultHostMetadataResolverFactory(t, func(c *Config) HostMetadataResolver {
+		factoryCalls.Add(1)
+		return func(ctx context.Context, host string) (*HostMetadata, error) {
+			return &HostMetadata{AccountID: testHMAccountID, WorkspaceID: testHMWorkspaceID}, nil
+		}
+	})
+
+	noopLoader := mockLoader(func(cfg *Config) error { return nil })
+	cfg := &Config{Host: testHMHost, Loaders: []Loader{noopLoader}}
+	require.NoError(t, cfg.EnsureResolved())
+
+	assert.Equal(t, int32(1), factoryCalls.Load(), "factory must be invoked exactly once per resolve")
+	assert.Equal(t, testHMAccountID, cfg.AccountID)
+	assert.Equal(t, testHMWorkspaceID, cfg.WorkspaceID)
+}
+
+func TestDefaultHostMetadataResolverFactory_PerConfigResolverTakesPrecedence(t *testing.T) {
+	var factoryCalls atomic.Int32
+	withDefaultHostMetadataResolverFactory(t, func(c *Config) HostMetadataResolver {
+		factoryCalls.Add(1)
+		return func(ctx context.Context, host string) (*HostMetadata, error) {
+			return &HostMetadata{AccountID: "factory-account"}, nil
+		}
+	})
+
+	noopLoader := mockLoader(func(cfg *Config) error { return nil })
+	cfg := &Config{
+		Host:    testHMHost,
+		Loaders: []Loader{noopLoader},
+		HostMetadataResolver: func(ctx context.Context, host string) (*HostMetadata, error) {
+			return &HostMetadata{AccountID: testHMAccountID}, nil
+		},
+	}
+	require.NoError(t, cfg.EnsureResolved())
+
+	assert.Equal(t, int32(0), factoryCalls.Load(), "factory must not be consulted when Config has its own resolver")
+	assert.Equal(t, testHMAccountID, cfg.AccountID)
+}
+
+func TestDefaultHostMetadataResolverFactory_NilResolverFromFactoryFallsThroughToHTTP(t *testing.T) {
+	withDefaultHostMetadataResolverFactory(t, func(c *Config) HostMetadataResolver {
+		return nil
+	})
+
+	noopLoader := mockLoader(func(cfg *Config) error { return nil })
+	cfg := &Config{
+		Host:    testHMHost,
+		Loaders: []Loader{noopLoader},
+		HTTPTransport: fixtures.SliceTransport{
+			{
+				Method:       "GET",
+				Resource:     "/.well-known/databricks-config",
+				ReuseRequest: true,
+				Status:       200,
+				Response:     `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `"}`,
+			},
+		},
+	}
+	require.NoError(t, cfg.EnsureResolved())
+
+	assert.Equal(t, testHMAccountID, cfg.AccountID)
+}

config/host_metadata.go

@@ -38,6 +38,20 @@ type HostMetadata struct {
 // This allows callers to provide cached metadata without the SDK making an HTTP call.
 type HostMetadataResolver func(ctx context.Context, host string) (*HostMetadata, error)
 
+// DefaultHostMetadataResolverFactory is consulted by [Config.EnsureResolved]
+// when [Config.HostMetadataResolver] is nil. When set, the factory is invoked
+// with the resolving Config and must return the resolver to use for that
+// Config (or nil to fall through to the SDK's default HTTP fetch).
+//
+// Intended for programs that want a single hook to install a caching or
+// otherwise-customised resolver across every Config they construct, without
+// per-site wiring. Set once from an init() block in a package that is
+// blank-imported by the main binary. Callers needing a per-Config resolver
+// should use [Config.HostMetadataResolver] instead.
+//
+// Experimental: subject to change.
+var DefaultHostMetadataResolverFactory func(*Config) HostMetadataResolver
+
 // getHostMetadata fetches the raw Databricks well-known configuration from
 // {host}/.well-known/databricks-config. The returned HostMetadata contains
 // raw values with no substitution (e.g., {account_id} placeholders are left