Tweak OAuth behavior

Author: stephanos

cliext/client.go

@@ -109,6 +109,9 @@ func (b *ClientOptionsBuilder) Build(ctx context.Context) (client.Options, error
 		profile.Namespace = cfg.Namespace
 	}
 
+	// Replace templated namespace in address, if present.
+	profile.Address = strings.ReplaceAll(profile.Address, "${namespace}", profile.Namespace)
+
 	// Set API key on profile if provided
 	if cfg.ApiKey != "" {
 		profile.APIKey = cfg.ApiKey
@@ -214,6 +217,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 but namespace is still the default to show a clearer error message.
+			if clientOpts.Namespace == "default" {
+				return client.Options{}, fmt.Errorf(
+					"Please specify a namespace with `--namespace` or configure one in your 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:   "localhost:17233",
+			Namespace: "default",
+		},
+	}
+
+	_, err := builder.Build(t.Context())
+
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "Please specify a namespace")
+}
+
 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")
+}