[Feature] Add support for Unified host without experimental flag

NEXT_CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Release v0.97.0
 
 ### New Features and Improvements
+* Add support for Unified host without experimental flag
 
 ### Bug Fixes
 

config/config.go

@@ -233,9 +233,6 @@ type Config struct {
 
 	// Keep track of the source of each attribute
 	attrSource map[string]Source
-
-	// Marker for unified hosts. Will be redundant once we can recognize unified hosts by their hostname.
-	Experimental_IsUnifiedHost bool `name:"experimental_is_unified_host" env:"DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST" auth:"-"`
 }
 
 // NewWithWorkspaceHost returns a new instance of the Config with the host set to
@@ -343,14 +340,11 @@ func (c *Config) IsAws() bool {
 }
 
 // IsAccountClient returns true if client is configured for Accounts API.
-// Panics if the config has the unified host flag set.
+// Note: This method does not support unified hosts. For unified hosts, use
+// HostType() or ConfigType() instead.
 //
 // Deprecated: Use HostType() if possible, or ConfigType() if necessary.
 func (c *Config) IsAccountClient() bool {
-	if c.Experimental_IsUnifiedHost {
-		panic("IsAccountClient cannot be used with unified hosts; use HostType() instead")
-	}
-
 	if c.AccountID != "" && c.isTesting {
 		return true
 	}
@@ -369,7 +363,7 @@ func (c *Config) IsAccountClient() bool {
 
 // HostType returns the type of host that the client is configured for.
 func (c *Config) HostType() HostType {
-	if c.Experimental_IsUnifiedHost {
+	if IsUnifiedHost(c.Host) {
 		return UnifiedHost
 	}
 

config/config_test.go

@@ -49,22 +49,12 @@ func TestHostType_AwsWorkspace(t *testing.T) {
 
 func TestHostType_Unified(t *testing.T) {
 	c := &Config{
-		Host:                       "https://unified.cloud.databricks.com",
-		AccountID:                  "123e4567-e89b-12d3-a456-426614174000",
-		Experimental_IsUnifiedHost: true,
+		Host:      "https://unified.databricks.com",
+		AccountID: "123e4567-e89b-12d3-a456-426614174000",
 	}
 	assert.Equal(t, UnifiedHost, c.HostType())
 }
 
-func TestIsAccountClient_PanicsOnUnifiedHost(t *testing.T) {
-	c := &Config{
-		Host:                       "https://unified.cloud.databricks.com",
-		AccountID:                  "test-account",
-		Experimental_IsUnifiedHost: true,
-	}
-	assert.Panics(t, func() { c.IsAccountClient() })
-}
-
 func TestNewWithWorkspaceHost(t *testing.T) {
 	c := &Config{
 		Host:               "https://accounts.cloud.databricks.com",
@@ -179,36 +169,35 @@ func TestConfig_getOidcEndpoints_unified(t *testing.T) {
 	}{
 		{
 			name:      "without trailing slash",
-			host:      "https://unified.cloud.databricks.com",
+			host:      "https://unified.databricks.com",
 			accountID: "abc",
 		},
 		{
 			name:      "with trailing slash",
-			host:      "https://unified.cloud.databricks.com/",
+			host:      "https://unified.databricks.com/",
 			accountID: "abc",
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := &Config{
-				Host:                       tt.host,
-				AccountID:                  tt.accountID,
-				Experimental_IsUnifiedHost: true,
+				Host:      tt.host,
+				AccountID: tt.accountID,
 				HTTPTransport: fixtures.SliceTransport{
 					{
 						Method:   "GET",
 						Resource: "/oidc/accounts/abc/.well-known/oauth-authorization-server",
 						Status:   200,
-						Response: `{"authorization_endpoint": "https://unified.cloud.databricks.com/oidc/accounts/abc/v1/authorize", "token_endpoint": "https://unified.cloud.databricks.com/oidc/accounts/abc/v1/token"}`,
+						Response: `{"authorization_endpoint": "https://unified.databricks.com/oidc/accounts/abc/v1/authorize", "token_endpoint": "https://unified.databricks.com/oidc/accounts/abc/v1/token"}`,
 					},
 				},
 			}
 			got, err := c.getOidcEndpoints(context.Background())
 			assert.NoError(t, err)
 			assert.Equal(t, &u2m.OAuthAuthorizationServer{
-				AuthorizationEndpoint: "https://unified.cloud.databricks.com/oidc/accounts/abc/v1/authorize",
-				TokenEndpoint:         "https://unified.cloud.databricks.com/oidc/accounts/abc/v1/token",
+				AuthorizationEndpoint: "https://unified.databricks.com/oidc/accounts/abc/v1/authorize",
+				TokenEndpoint:         "https://unified.databricks.com/oidc/accounts/abc/v1/token",
 			}, got)
 		})
 	}
@@ -285,28 +274,27 @@ func TestConfig_getOAuthArgument_Unified(t *testing.T) {
 	}{
 		{
 			name:      "without trailing slash",
-			host:      "https://unified.cloud.databricks.com",
+			host:      "https://unified.databricks.com",
 			accountID: "account-123",
 		},
 		{
 			name:      "with trailing slash",
-			host:      "https://unified.cloud.databricks.com/",
+			host:      "https://unified.databricks.com/",
 			accountID: "account-123",
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := &Config{
-				Host:                       tt.host,
-				AccountID:                  tt.accountID,
-				Experimental_IsUnifiedHost: true,
+				Host:      tt.host,
+				AccountID: tt.accountID,
 			}
 			rawGot, err := c.getOAuthArgument()
 			assert.NoError(t, err)
 			got, ok := rawGot.(u2m.UnifiedOAuthArgument)
 			assert.True(t, ok, "Expected UnifiedOAuthArgument")
-			assert.Equal(t, "https://unified.cloud.databricks.com", got.GetHost())
+			assert.Equal(t, "https://unified.databricks.com", got.GetHost())
 			assert.Equal(t, "account-123", got.GetAccountId())
 		})
 	}

config/spog.go

@@ -0,0 +1,85 @@
+package config
+
+import (
+	"net/url"
+	"strings"
+)
+
+// IsUnifiedHost returns true if the given host is a unified host that supports
+// both workspace-level and account-level APIs. This matches the SPOG domain checker
+// logic from the platform.
+//
+// A hostname is considered a unified host if it matches one of the SPOG URL patterns:
+// - *.databricks.com (without subdomains, e.g., company.databricks.com)
+// - *.azuredatabricks.net (without subdomains, e.g., company.azuredatabricks.net)
+func IsUnifiedHost(host string) bool {
+	if host == "" {
+		return false
+	}
+
+	hostname := extractHostname(host)
+	if hostname == "" {
+		return false
+	}
+
+	return matchesSpogUrlPattern(hostname)
+}
+
+// extractHostname parses a URL or hostname string and returns the hostname component.
+// It handles both full URLs (with scheme) and bare hostnames.
+func extractHostname(host string) string {
+	// Parse the URL to extract just the hostname
+	parsedHost, err := url.Parse(host)
+	if err != nil {
+		return ""
+	}
+
+	// If no host was parsed, assume the scheme wasn't included
+	hostname := parsedHost.Hostname()
+	if hostname == "" {
+		parsedHost, err = url.Parse("https://" + host)
+		if err != nil {
+			return ""
+		}
+		hostname = parsedHost.Hostname()
+	}
+
+	return hostname
+}
+
+// matchesSpogUrlPattern checks if the hostname matches any SPOG URL pattern.
+// For production (which is what the SDK primarily targets), this includes:
+// - *.databricks.com
+// - *.azuredatabricks.net
+func matchesSpogUrlPattern(hostname string) bool {
+	// SPOG URL patterns for production
+	spogUrlSuffixes := []string{".databricks.com", ".azuredatabricks.net"}
+
+	// Check if hostname matches any SPOG URL pattern
+	for _, suffix := range spogUrlSuffixes {
+		if hostnameMatchesSpogUrlSuffix(hostname, suffix) {
+			return true
+		}
+	}
+
+	return false
+}
+
+// hostnameMatchesSpogUrlSuffix checks if hostname ends with the given suffix
+// and that the prefix (before the suffix) doesn't contain any dots.
+// This ensures that only single-level subdomains are matched.
+//
+// Examples:
+//   - "company.databricks.com" with suffix ".databricks.com" -> true (prefix "company" has no dots)
+//   - "dbc-12345.cloud.databricks.com" with suffix ".databricks.com" -> false (prefix "dbc-12345.cloud" has dots)
+func hostnameMatchesSpogUrlSuffix(hostname string, suffix string) bool {
+	if !strings.HasSuffix(hostname, suffix) {
+		return false
+	}
+
+	// Extract the prefix before the suffix
+	prefix := strings.TrimSuffix(hostname, suffix)
+
+	// Check that the prefix doesn't contain any dots (no subdomains)
+	return !strings.Contains(prefix, ".")
+}