Migrating custom action client for client v2

pkg/azure/armauth/auth.go

@@ -29,8 +29,8 @@ type ArmConfig struct {
 	// TODO: Migrate authenticator and clients to new azure sdk - https://github.com/project-radius/radius/issues/4268
 	Auth autorest.Authorizer
 
-	// ClientOption is the client v2 option including new client credentials.
-	ClientOption clientv2.AzureClientOption
+	// ClientOptions is the client v2 options including new client credentials.
+	ClientOptions clientv2.ClientOptions
 }
 
 // GetArmConfig gets the configuration we use for managing ARM resources
@@ -47,8 +47,8 @@ func GetArmConfig() (*ArmConfig, error) {
 	}
 
 	return &ArmConfig{
-		Auth:         auth,
-		ClientOption: clientv2.AzureClientOption{Cred: cred},
+		Auth:          auth,
+		ClientOptions: clientv2.ClientOptions{Cred: cred},
 	}, nil
 }
 

pkg/azure/clientv2/clients.go

@@ -20,20 +20,29 @@ var defaultClientOptions = &arm.ClientOptions{
 	},
 }
 
-// AzureClientOption represents the client option for azure sdk client including authentication.
-type AzureClientOption struct {
+// ClientOptions represents the client option for azure sdk client including authentication.
+type ClientOptions struct {
 	// Cred represents a credential for OAuth token.
 	Cred azcore.TokenCredential
+
+	// BaseURI is the base URI for the client to be created.
+	BaseURI string
+}
+
+// SetBaseURI updates the base URI for the Azure Client Options.
+func (options *ClientOptions) SetBaseURI(baseURI string) *ClientOptions {
+	options.BaseURI = baseURI
+	return options
 }
 
 // NewFederatedIdentityClient creates new federated identity client.
-func NewFederatedIdentityClient(subscriptionID string, option *AzureClientOption) (*armmsi.FederatedIdentityCredentialsClient, error) {
+func NewFederatedIdentityClient(subscriptionID string, options *ClientOptions) (*armmsi.FederatedIdentityCredentialsClient, error) {
 	// TODO: Add LRU cache to maintain the clients.
-	return armmsi.NewFederatedIdentityCredentialsClient(subscriptionID, option.Cred, defaultClientOptions)
+	return armmsi.NewFederatedIdentityCredentialsClient(subscriptionID, options.Cred, defaultClientOptions)
 }
 
 // NewUserAssignedIdentityClient creates new user assigned managed identity client.
-func NewUserAssignedIdentityClient(subscriptionID string, option *AzureClientOption) (*armmsi.UserAssignedIdentitiesClient, error) {
+func NewUserAssignedIdentityClient(subscriptionID string, options *ClientOptions) (*armmsi.UserAssignedIdentitiesClient, error) {
 	// TODO: Add LRU cache to maintain the clients.
-	return armmsi.NewUserAssignedIdentitiesClient(subscriptionID, option.Cred, defaultClientOptions)
+	return armmsi.NewUserAssignedIdentitiesClient(subscriptionID, options.Cred, defaultClientOptions)
 }

pkg/azure/clientv2/common.go

@@ -9,18 +9,13 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources"
 )
 
 // GetResourceGroupLocation returns the location of the resource group.
-func GetResourceGroupLocation(ctx context.Context, credential azcore.TokenCredential, subscriptionID string, resourceGroupName string) (*string, error) {
-	client, err := armresources.NewResourceGroupsClient(subscriptionID, credential, &arm.ClientOptions{})
-	if err != nil {
-		return nil, err
-	}
-
+func GetResourceGroupLocation(ctx context.Context, subscriptionID string, resourceGroupName string, options *ClientOptions) (*string, error) {
+	client, err := armresources.NewResourceGroupsClient(subscriptionID, options.Cred, &arm.ClientOptions{})
 	if err != nil {
 		return nil, err
 	}

pkg/azure/clientv2/customaction.go

@@ -7,6 +7,7 @@ package clientv2
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"net/http"
 	"net/url"
@@ -22,53 +23,58 @@ import (
 	"github.com/project-radius/radius/pkg/ucp/resources"
 )
 
+// CustomActionClient is the client to invoke custom actions on Azure resources.
+// Ex: listSecrets on a MongoDatabase.
 type CustomActionClient struct {
-	*BaseClient
+	Client   *armresources.Client
+	Pipeline *runtime.Pipeline
+	BaseURI  string
 }
 
 // NewCustomActionClient creates an instance of the CustomActionClient with the default Base URI.
-func NewCustomActionClient(subscriptionID string, credential azcore.TokenCredential) (*BaseClient, error) {
-	client, err := NewCustomActionClientWithBaseURI(DefaultBaseURI, subscriptionID, credential)
-	if err != nil {
-		return nil, err
+func NewCustomActionClient(subscriptionID string, options *ClientOptions) (*CustomActionClient, error) {
+	baseURI := DefaultBaseURI
+	if options.BaseURI != "" {
+		baseURI = options.BaseURI
 	}
 
-	return client, err
-}
-
-// NewCustomActionClientWithBaseURI creates an instance of the CustomActionClient with a Base URI.
-func NewCustomActionClientWithBaseURI(baseURI string, subscriptionID string, credential azcore.TokenCredential) (*BaseClient, error) {
-	options := &arm.ClientOptions{
+	armClientOptions := &arm.ClientOptions{
 		ClientOptions: azcore.ClientOptions{
 			Cloud: cloud.Configuration{
 				Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
 					cloud.ResourceManager: {
+						Audience: "https://management.core.windows.net",
 						Endpoint: baseURI,
 					},
 				},
 			},
 		},
 	}
 
-	client, err := armresources.NewClient(subscriptionID, credential, options)
+	client, err := armresources.NewClient(subscriptionID, options.Cred, armClientOptions)
 	if err != nil {
 		return nil, err
 	}
 
-	pipeline, err := armruntime.NewPipeline(moduleName, moduleVersion, credential, runtime.PipelineOptions{}, options)
+	pipeline, err := armruntime.NewPipeline(ModuleName, ModuleVersion, options.Cred, runtime.PipelineOptions{}, armClientOptions)
 	if err != nil {
 		return nil, err
 	}
 
-	return &BaseClient{
+	return &CustomActionClient{
 		Client:   client,
 		Pipeline: &pipeline,
 		BaseURI:  baseURI,
 	}, nil
 }
 
 type ClientCustomActionResponse struct {
-	armresources.GenericResource
+	CustomActionResponseBody
+}
+
+type CustomActionResponseBody struct {
+	Body     map[string]any
+	Response http.Response
 }
 
 type ClientBeginCustomActionOptions struct {
@@ -77,32 +83,25 @@ type ClientBeginCustomActionOptions struct {
 	apiVersion string
 }
 
-// New creates an instance of the CustomActionClient with the default Base URI.
-func NewCustomActionRequestOptions(resourceID, action, apiVersion string) *ClientBeginCustomActionOptions {
-	// FIXME: This is to validate the resourceID.
-	_, err := resources.ParseResource(resourceID)
-	if err != nil {
-		return nil
-	}
-
+// NewClientCustomActionOptions creates an instance of the CustomActionClientOptions.
+func NewClientBeginCustomActionOptions(resourceID, action, apiVersion string) *ClientBeginCustomActionOptions {
 	return &ClientBeginCustomActionOptions{
 		resourceID: resourceID,
 		action:     action,
 		apiVersion: apiVersion,
 	}
 }
 
-func (client *CustomActionClient) BeginCustomAction(ctx context.Context, opts *ClientBeginCustomActionOptions) (*runtime.Poller[ClientCustomActionResponse], error) {
+func (client *CustomActionClient) BeginCustomAction(ctx context.Context, opts *ClientBeginCustomActionOptions) (*ClientCustomActionResponse, error) {
 	resp, err := client.customAction(ctx, opts)
 	if err != nil {
 		return nil, err
 	}
 
-	// FIXME: Is this the right way?
-	return runtime.NewPoller[ClientCustomActionResponse](resp, *client.Pipeline, nil)
+	return resp, nil
 }
 
-func (client *CustomActionClient) customAction(ctx context.Context, opts *ClientBeginCustomActionOptions) (*http.Response, error) {
+func (client *CustomActionClient) customAction(ctx context.Context, opts *ClientBeginCustomActionOptions) (*ClientCustomActionResponse, error) {
 	req, err := client.customActionCreateRequest(ctx, opts)
 	if err != nil {
 		return nil, err
@@ -112,14 +111,33 @@ func (client *CustomActionClient) customAction(ctx context.Context, opts *Client
 	if err != nil {
 		return nil, err
 	}
-	if !runtime.HasStatusCode(resp, http.StatusAccepted, http.StatusNoContent) {
+
+	if !runtime.HasStatusCode(resp, http.StatusOK, http.StatusAccepted, http.StatusNoContent) {
 		return nil, runtime.NewResponseError(resp)
 	}
-	return resp, nil
+
+	body := map[string]any{}
+	err = json.NewDecoder(resp.Body).Decode(&body)
+	if err != nil {
+		return nil, err
+	}
+
+	return &ClientCustomActionResponse{
+		CustomActionResponseBody{
+			Body:     body,
+			Response: *resp,
+		},
+	}, nil
 }
 
 func (client *CustomActionClient) customActionCreateRequest(ctx context.Context, opts *ClientBeginCustomActionOptions) (*policy.Request, error) {
-	urlPath := "/{resourceID}/{action}"
+	_, err := resources.ParseResource(opts.resourceID)
+	if err != nil {
+		return nil, err
+	}
+
+	urlPath := "{resourceID}/{action}"
+
 	if opts.resourceID == "" {
 		return nil, errors.New("resourceID cannot be empty")
 	}
@@ -130,7 +148,6 @@ func (client *CustomActionClient) customActionCreateRequest(ctx context.Context,
 	}
 	urlPath = strings.ReplaceAll(urlPath, "{action}", url.PathEscape(opts.action))
 
-	// FIXME: Is joining BaseURI and URLPath going to give us a wrong URL?
 	req, err := runtime.NewRequest(ctx, http.MethodPost, runtime.JoinPaths(client.BaseURI, urlPath))
 	if err != nil {
 		return nil, err

pkg/azure/clientv2/resourcedeploymentclient.go

@@ -111,6 +111,7 @@ func NewDeploymentsClientWithBaseURI(credential azcore.TokenCredential, subscrip
 			Cloud: cloud.Configuration{
 				Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
 					cloud.ResourceManager: {
+						Audience: "https://management.core.windows.net",
 						Endpoint: baseURI,
 					},
 				},
@@ -122,7 +123,7 @@ func NewDeploymentsClientWithBaseURI(credential azcore.TokenCredential, subscrip
 		return nil, err
 	}
 
-	pipeline, err := armruntime.NewPipeline(moduleName, moduleVersion, credential, runtime.PipelineOptions{}, options)
+	pipeline, err := armruntime.NewPipeline(ModuleName, ModuleVersion, credential, runtime.PipelineOptions{}, options)
 	if err != nil {
 		return nil, err
 	}
@@ -135,22 +136,19 @@ func NewDeploymentsClientWithBaseURI(credential azcore.TokenCredential, subscrip
 }
 
 type ClientBeginCreateOrUpdateOptions struct {
-	resourceID  string
-	resumeToken string
-	apiVersion  string
+	resourceID string
+	apiVersion string
 }
 
-func NewClientBeginCreateOrUpdateOptions(resourceID, resumeToken, apiVersion string) *ClientBeginCreateOrUpdateOptions {
-	// FIXME: This is to validate the resourceID.
+func NewClientBeginCreateOrUpdateOptions(resourceID, apiVersion string) *ClientBeginCreateOrUpdateOptions {
 	_, err := resources.ParseResource(resourceID)
 	if err != nil {
 		return nil
 	}
 
 	return &ClientBeginCreateOrUpdateOptions{
-		resourceID:  resourceID,
-		resumeToken: resumeToken,
-		apiVersion:  apiVersion,
+		resourceID: resourceID,
+		apiVersion: apiVersion,
 	}
 }
 

pkg/azure/clientv2/resourcedeploymentoperationsclient.go

@@ -39,6 +39,7 @@ func NewResourceDeploymentOperationsClientWithBaseURI(cred azcore.TokenCredentia
 			Cloud: cloud.Configuration{
 				Services: map[cloud.ServiceName]cloud.ServiceConfiguration{
 					cloud.ResourceManager: {
+						Audience: "https://management.core.windows.net",
 						Endpoint: baseURI,
 					},
 				},

pkg/azure/clientv2/types.go

@@ -5,21 +5,8 @@
 
 package clientv2
 
-import (
-	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
-	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources"
-)
-
 const (
 	DefaultBaseURI = "https://management.azure.com"
-
-	// FIXM: Any ideas for moduleName and moduleVersion?
-	moduleName    = "radius"
-	moduleVersion = "public-preview"
+	ModuleName     = "radius"
+	ModuleVersion  = "public-preview"
 )
-
-type BaseClient struct {
-	Client   *armresources.Client
-	Pipeline *runtime.Pipeline
-	BaseURI  string
-}

pkg/corerp/handlers/azure_federatedidentity.go

@@ -129,7 +129,7 @@ func (handler *azureFederatedIdentityHandler) Put(ctx context.Context, options *
 	subID := rID.FindScope(resources.SubscriptionsSegment)
 	rgName := rID.FindScope(resources.ResourceGroupsSegment)
 
-	client, err := clientv2.NewFederatedIdentityClient(subID, &handler.arm.ClientOption)
+	client, err := clientv2.NewFederatedIdentityClient(subID, &handler.arm.ClientOptions)
 	if err != nil {
 		return nil, err
 	}
@@ -179,7 +179,9 @@ func (handler *azureFederatedIdentityHandler) Delete(ctx context.Context, option
 		return err
 	}
 
-	client, err := clientv2.NewFederatedIdentityClient(rID.FindScope(resources.SubscriptionsSegment), &handler.arm.ClientOption)
+	subscriptionID := rID.FindScope(resources.SubscriptionsSegment)
+
+	client, err := clientv2.NewFederatedIdentityClient(subscriptionID, &handler.arm.ClientOptions)
 	if err != nil {
 		return err
 	}

pkg/corerp/handlers/azure_userassigned_managedidentity.go

@@ -13,7 +13,6 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/msi/armmsi"
 	"github.com/Azure/go-autorest/autorest/to"
 	"github.com/project-radius/radius/pkg/azure/armauth"
-	"github.com/project-radius/radius/pkg/azure/clients"
 	"github.com/project-radius/radius/pkg/azure/clientv2"
 	"github.com/project-radius/radius/pkg/radlogger"
 	"github.com/project-radius/radius/pkg/resourcemodel"
@@ -64,7 +63,7 @@ func (handler *azureUserAssignedManagedIdentityHandler) Put(ctx context.Context,
 		return nil, err
 	}
 
-	rgLocation, err := clients.GetResourceGroupLocation(ctx, *handler.arm, subID, rgName)
+	rgLocation, err := clientv2.GetResourceGroupLocation(ctx, subID, rgName, &handler.arm.ClientOptions)
 	if err != nil {
 		return properties, err
 	}
@@ -78,7 +77,7 @@ func (handler *azureUserAssignedManagedIdentityHandler) Put(ctx context.Context,
 		return nil, fmt.Errorf("azure federated identity does not support %s region now. unsupported regions: %q", resourceLocation, federatedUnsupportedRegions)
 	}
 
-	msiClient, err := clientv2.NewUserAssignedIdentityClient(subID, &handler.arm.ClientOption)
+	msiClient, err := clientv2.NewUserAssignedIdentityClient(subID, &handler.arm.ClientOptions)
 	if err != nil {
 		return nil, err
 	}
@@ -93,7 +92,7 @@ func (handler *azureUserAssignedManagedIdentityHandler) Put(ctx context.Context,
 	properties[UserAssignedIdentityClientIDKey] = to.String(identity.Properties.ClientID)
 	properties[UserAssignedIdentityTenantIDKey] = to.String(identity.Properties.TenantID)
 
-	options.Resource.Identity = resourcemodel.NewARMIdentity(&options.Resource.ResourceType, properties[UserAssignedIdentityIDKey], clients.GetAPIVersionFromUserAgent(msi.UserAgent()))
+	options.Resource.Identity = resourcemodel.NewARMIdentity(&options.Resource.ResourceType, properties[UserAssignedIdentityIDKey], clientv2.GetAPIVersionFromUserAgent(msi.UserAgent()))
 	logger.WithValues(
 		radlogger.LogFieldResourceID, *identity.ID,
 		radlogger.LogFieldLocalID, outputresource.LocalIDUserAssignedManagedIdentity).Info("Created managed identity for KeyVault access")
@@ -112,7 +111,9 @@ func (handler *azureUserAssignedManagedIdentityHandler) Delete(ctx context.Conte
 		return err
 	}
 
-	msiClient, err := clientv2.NewUserAssignedIdentityClient(parsed.FindScope(resources.SubscriptionsSegment), &handler.arm.ClientOption)
+	subscriptionID := parsed.FindScope(resources.SubscriptionsSegment)
+
+	msiClient, err := clientv2.NewUserAssignedIdentityClient(subscriptionID, &handler.arm.ClientOptions)
 	if err != nil {
 		return err
 	}