Tweak OAuth behavior (#913)

<!--- Note to EXTERNAL Contributors -->
<!-- Thanks for opening a PR! 
If it is a significant code change, please **make sure there is an open
issue** for this.
We work best with you when we have accepted the idea first before you
code. -->

<!--- For ALL Contributors 👇 -->

## What was changed
<!-- Describe what has changed in this PR -->

(1) Ability to update the `address` env profile field when storing OAuth
config.
(2) Ability to use `${namespace}` in `address` to refer to namespace at
runtime.
(3) Show helpful error message to user in case they have OAuth but still
use `default` namespace.

## Checklist
<!--- add/delete as needed --->

1. Closes <!-- add issue number here -->

2. How was this tested:
<!--- Please describe how you tested your changes/how we can test them
-->

3. Any docs updates needed?
<!--- update README if applicable
      or point out where to update docs.temporal.io -->

Author: stephanos

cliext/client.go

@@ -14,6 +14,9 @@ import (
 	"google.golang.org/grpc"
 )
 
+// addressNamespaceTemplate is the placeholder in addresses that gets replaced with the namespace.
+const addressNamespaceTemplate = "${namespace}"
+
 // ClientOptionsBuilder contains options for building SDK client.Options.
 type ClientOptionsBuilder struct {
 	// CommonOptions contains common CLI options including profile config.
@@ -103,12 +106,22 @@ func (b *ClientOptionsBuilder) Build(ctx context.Context) (client.Options, error
 	} else if profile.Address == "" {
 		profile.Address = cfg.Address
 	}
+	var namespaceExplicitlySet bool
 	if cfg.FlagSet != nil && cfg.FlagSet.Changed("namespace") {
+		namespaceExplicitlySet = true
 		profile.Namespace = cfg.Namespace
-	} else if profile.Namespace == "" {
+	} else if profile.Namespace != "" {
+		namespaceExplicitlySet = true
+	} else {
 		profile.Namespace = cfg.Namespace
 	}
 
+	// Replace templated namespace in address, if present.
+	addressHasNamespaceTemplate := strings.Contains(profile.Address, addressNamespaceTemplate)
+	if addressHasNamespaceTemplate {
+		profile.Address = strings.ReplaceAll(profile.Address, addressNamespaceTemplate, profile.Namespace)
+	}
+
 	// Set API key on profile if provided
 	if cfg.ApiKey != "" {
 		profile.APIKey = cfg.ApiKey
@@ -214,6 +227,11 @@ func (b *ClientOptionsBuilder) Build(ctx context.Context) (client.Options, error
 		}
 		// Only set credentials if OAuth is configured with an access token
 		if result.OAuth != nil && result.OAuth.Token != nil && result.OAuth.Token.AccessToken != "" {
+			// Error if OAuth is configured with templated address but namespace was not explicitly set.
+			if addressHasNamespaceTemplate && !namespaceExplicitlySet {
+				return client.Options{}, fmt.Errorf(
+					"namespace is required to be set via `--namespace` or configured in profile.")
+			}
 			creds := &oauthCredentials{
 				builder:        b,
 				config:         result.OAuth,

cliext/client_test.go

@@ -80,6 +80,23 @@ func TestClientOptionsBuilder_OptionsPassedThrough(t *testing.T) {
 	assert.NotNil(t, opts.HeadersProvider)
 }
 
+func TestClientOptionsBuilder_NamespaceReplacementInAddress(t *testing.T) {
+	builder := &cliext.ClientOptionsBuilder{
+		CommonOptions: cliext.CommonOptions{
+			DisableConfigFile: true,
+		},
+		ClientOptions: cliext.ClientOptions{
+			Address:   "${namespace}.api.temporal.io:7233",
+			Namespace: "my-namespace",
+		},
+	}
+
+	opts, err := builder.Build(t.Context())
+
+	require.NoError(t, err)
+	assert.Equal(t, "my-namespace.api.temporal.io:7233", opts.HostPort)
+}
+
 func TestClientOptionsBuilder_OAuth_ValidToken(t *testing.T) {
 	s := newMockOAuthServer(t)
 	configFile := createTestOAuthConfig(t, s.tokenURL, time.Now().Add(time.Hour))
@@ -89,7 +106,8 @@ func TestClientOptionsBuilder_OAuth_ValidToken(t *testing.T) {
 			ConfigFile: configFile,
 		},
 		ClientOptions: cliext.ClientOptions{
-			Address: "localhost:17233",
+			Address:   "localhost:17233",
+			Namespace: "my-namespace",
 		},
 	}
 	opts, err := builder.Build(t.Context())
@@ -111,7 +129,8 @@ func TestClientOptionsBuilder_OAuth_Refresh(t *testing.T) {
 			ConfigFile: configFile,
 		},
 		ClientOptions: cliext.ClientOptions{
-			Address: "localhost:17233",
+			Address:   "localhost:17233",
+			Namespace: "my-namespace",
 		},
 	}
 	opts, err := builder.Build(context.Background())
@@ -141,7 +160,8 @@ func TestClientOptionsBuilder_OAuth_NoConfigFile(t *testing.T) {
 			DisableConfigFile: true,
 		},
 		ClientOptions: cliext.ClientOptions{
-			Address: "localhost:17233",
+			Address:   "localhost:17233",
+			Namespace: "my-namespace",
 		},
 	}
 	opts, err := builder.Build(t.Context())
@@ -160,7 +180,8 @@ func TestClientOptionsBuilder_OAuth_NoOAuthConfigured(t *testing.T) {
 			ConfigFile: configFile,
 		},
 		ClientOptions: cliext.ClientOptions{
-			Address: "localhost:17233",
+			Address:   "localhost:17233",
+			Namespace: "my-namespace",
 		},
 	}
 	opts, err := builder.Build(t.Context())
@@ -179,7 +200,8 @@ func TestClientOptionsBuilder_OAuth_DisableConfigFile(t *testing.T) {
 			DisableConfigFile: true, // disables loading config file
 		},
 		ClientOptions: cliext.ClientOptions{
-			Address: "localhost:17233",
+			Address:   "localhost:17233",
+			Namespace: "my-namespace",
 		},
 	}
 	opts, err := builder.Build(t.Context())
@@ -197,8 +219,9 @@ func TestClientOptionsBuilder_OAuth_APIKeyTakesPrecedence(t *testing.T) {
 			ConfigFile: configFileWithOAuth,
 		},
 		ClientOptions: cliext.ClientOptions{
-			Address: "localhost:17233",
-			ApiKey:  "explicit-api-key", // takes precedence
+			Address:   "localhost:17233",
+			Namespace: "my-namespace",
+			ApiKey:    "explicit-api-key", // takes precedence
 		},
 	}
 	opts, err := builder.Build(t.Context())
@@ -207,6 +230,25 @@ func TestClientOptionsBuilder_OAuth_APIKeyTakesPrecedence(t *testing.T) {
 	assert.NotNil(t, opts.Credentials)
 }
 
+func TestClientOptionsBuilder_OAuth_NonDefaultNamespaceRequired(t *testing.T) {
+	configFileWithOAuth := createTestOAuthConfig(t, "", time.Now().Add(time.Hour))
+
+	builder := &cliext.ClientOptionsBuilder{
+		CommonOptions: cliext.CommonOptions{
+			ConfigFile: configFileWithOAuth,
+		},
+		ClientOptions: cliext.ClientOptions{
+			Address:   "${namespace}.api.temporal.io:7233",
+			Namespace: "default",
+		},
+	}
+
+	_, err := builder.Build(t.Context())
+
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "namespace is required")
+}
+
 func createTestOAuthConfig(t *testing.T, tokenURL string, expiry time.Time) string {
 	t.Helper()
 	configFile := filepath.Join(t.TempDir(), "config.toml")

cliext/config.oauth.go

@@ -169,6 +169,8 @@ type StoreClientOAuthOptions struct {
 	ProfileName string
 	// OAuth is the OAuth configuration to store. If nil, removes OAuth from the profile.
 	OAuth *OAuthConfig
+	// Address is the server address to store in the profile. If empty, address is not modified.
+	Address string
 	// EnvLookup overrides environment variable lookup. If nil, uses os.LookupEnv.
 	EnvLookup envconfig.EnvLookup
 }
@@ -215,6 +217,21 @@ func StoreClientOAuth(opts StoreClientOAuthOptions) error {
 		return err
 	}
 
+	// Set address if provided.
+	if opts.Address != "" {
+		profileSection, ok := existingRaw["profile"].(map[string]any)
+		if !ok {
+			profileSection = make(map[string]any)
+			existingRaw["profile"] = profileSection
+		}
+		profile, ok := profileSection[profileName].(map[string]any)
+		if !ok {
+			profile = make(map[string]any)
+			profileSection[profileName] = profile
+		}
+		profile["address"] = opts.Address
+	}
+
 	// Marshal back to TOML.
 	var buf bytes.Buffer
 	enc := toml.NewEncoder(&buf)

cliext/config.oauth_test.go

@@ -258,3 +258,31 @@ address = "other:7233"
 	assert.Contains(t, string(content), "other:7233")
 	assert.Contains(t, string(content), "test-client")
 }
+
+func TestStoreClientOAuth_WithAddress(t *testing.T) {
+	f, err := os.CreateTemp("", "temporal-config-*.toml")
+	require.NoError(t, err)
+	defer os.Remove(f.Name())
+	require.NoError(t, f.Close())
+
+	// Store OAuth config with address
+	err = cliext.StoreClientOAuth(cliext.StoreClientOAuthOptions{
+		ConfigFilePath: f.Name(),
+		OAuth: &cliext.OAuthConfig{
+			ClientConfig: &oauth2.Config{
+				ClientID: "test-client",
+			},
+			Token: &oauth2.Token{
+				AccessToken: "test-token",
+			},
+		},
+		Address: "${namespace}.api.temporal.io:7233",
+	})
+	require.NoError(t, err)
+
+	// Verify address was written
+	content, err := os.ReadFile(f.Name())
+	require.NoError(t, err)
+	assert.Contains(t, string(content), `address = "${namespace}.api.temporal.io:7233"`)
+	assert.Contains(t, string(content), "test-client")
+}