cosmos - migrate service to go-azure-sdk

internal/services/cosmos/client/client.go

@@ -6,11 +6,12 @@ package client
 import (
 	"fmt"
 
-	"github.com/Azure/azure-sdk-for-go/services/cosmos-db/mgmt/2021-10-15/documentdb" // nolint: staticcheck
 	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2022-05-15/sqldedicatedgateway"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2022-11-15/mongorbacs"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2023-04-15/managedcassandras"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/cosmosdb"
+	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/rbacs"
+	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/restorables"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/postgresqlhsc/2022-11-08/clusters"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/postgresqlhsc/2022-11-08/configurations"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/postgresqlhsc/2022-11-08/firewallrules"
@@ -25,19 +26,10 @@ type Client struct {
 	FirewallRulesClient       *firewallrules.FirewallRulesClient
 	ManagedCassandraClient    *managedcassandras.ManagedCassandrasClient
 	MongoRBACClient           *mongorbacs.MongorbacsClient
+	RbacsClient               *rbacs.RbacsClient
+	RestorablesClient         *restorables.RestorablesClient
 	RolesClient               *roles.RolesClient
 	SqlDedicatedGatewayClient *sqldedicatedgateway.SqlDedicatedGatewayClient
-
-	// Track 1
-	CassandraClient                  *documentdb.CassandraResourcesClient
-	DatabaseClient                   *documentdb.DatabaseAccountsClient
-	GremlinClient                    *documentdb.GremlinResourcesClient
-	MongoDbClient                    *documentdb.MongoDBResourcesClient
-	NotebookWorkspaceClient          *documentdb.NotebookWorkspacesClient
-	RestorableDatabaseAccountsClient *documentdb.RestorableDatabaseAccountsClient
-	SqlClient                        *documentdb.SQLResourcesClient
-	SqlResourceClient                *documentdb.SQLResourcesClient
-	TableClient                      *documentdb.TableResourcesClient
 }
 
 func NewClient(o *common.ClientOptions) (*Client, error) {
@@ -77,65 +69,40 @@ func NewClient(o *common.ClientOptions) (*Client, error) {
 	}
 	o.Configure(mongorbacsClient.Client, o.Authorizers.ResourceManager)
 
+	rbacsClient, err := rbacs.NewRbacsClientWithBaseURI(o.Environment.ResourceManager)
+	if err != nil {
+		return nil, fmt.Errorf("building RBACs client: %+v", err)
+	}
+	o.Configure(rbacsClient.Client, o.Authorizers.ResourceManager)
+
 	rolesClient, err := roles.NewRolesClientWithBaseURI(o.Environment.ResourceManager)
 	if err != nil {
 		return nil, fmt.Errorf("building Roles client: %+v", err)
 	}
 	o.Configure(rolesClient.Client, o.Authorizers.ResourceManager)
 
+	restorablesClient, err := restorables.NewRestorablesClientWithBaseURI(o.Environment.ResourceManager)
+	if err != nil {
+		return nil, fmt.Errorf("building Restorables client: %+v", err)
+	}
+	o.Configure(restorablesClient.Client, o.Authorizers.ResourceManager)
+
 	sqlDedicatedGatewayClient, err := sqldedicatedgateway.NewSqlDedicatedGatewayClientWithBaseURI(o.Environment.ResourceManager)
 	if err != nil {
 		return nil, fmt.Errorf("building Sql Dedicated Gateway client: %+v", err)
 	}
 	o.Configure(sqlDedicatedGatewayClient.Client, o.Authorizers.ResourceManager)
 
-	// Track 1
-	cassandraClient := documentdb.NewCassandraResourcesClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&cassandraClient.Client, o.ResourceManagerAuthorizer)
-
-	databaseClient := documentdb.NewDatabaseAccountsClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&databaseClient.Client, o.ResourceManagerAuthorizer)
-
-	gremlinClient := documentdb.NewGremlinResourcesClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&gremlinClient.Client, o.ResourceManagerAuthorizer)
-
-	mongoDbClient := documentdb.NewMongoDBResourcesClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&mongoDbClient.Client, o.ResourceManagerAuthorizer)
-
-	notebookWorkspaceClient := documentdb.NewNotebookWorkspacesClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&notebookWorkspaceClient.Client, o.ResourceManagerAuthorizer)
-
-	restorableDatabaseAccountsClient := documentdb.NewRestorableDatabaseAccountsClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&restorableDatabaseAccountsClient.Client, o.ResourceManagerAuthorizer)
-
-	sqlClient := documentdb.NewSQLResourcesClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&sqlClient.Client, o.ResourceManagerAuthorizer)
-
-	sqlResourceClient := documentdb.NewSQLResourcesClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&sqlResourceClient.Client, o.ResourceManagerAuthorizer)
-
-	tableClient := documentdb.NewTableResourcesClientWithBaseURI(o.ResourceManagerEndpoint, o.SubscriptionId)
-	o.ConfigureClient(&tableClient.Client, o.ResourceManagerAuthorizer)
-
 	return &Client{
 		ManagedCassandraClient:    managedCassandraClient,
 		ClustersClient:            clustersClient,
 		ConfigurationsClient:      configurationsClient,
 		CosmosDBClient:            cosmosdbClient,
 		FirewallRulesClient:       firewallRulesClient,
 		MongoRBACClient:           mongorbacsClient,
+		RbacsClient:               rbacsClient,
+		RestorablesClient:         restorablesClient,
 		RolesClient:               rolesClient,
 		SqlDedicatedGatewayClient: sqlDedicatedGatewayClient,
-
-		// Track 1
-		CassandraClient:                  &cassandraClient,
-		DatabaseClient:                   &databaseClient,
-		GremlinClient:                    &gremlinClient,
-		MongoDbClient:                    &mongoDbClient,
-		NotebookWorkspaceClient:          &notebookWorkspaceClient,
-		RestorableDatabaseAccountsClient: &restorableDatabaseAccountsClient,
-		SqlClient:                        &sqlClient,
-		SqlResourceClient:                &sqlResourceClient,
-		TableClient:                      &tableClient,
 	}, nil
 }

internal/services/cosmos/common/autoscale_settings.go

@@ -6,29 +6,11 @@ package common
 import (
 	"log"
 
-	"github.com/Azure/azure-sdk-for-go/services/cosmos-db/mgmt/2021-10-15/documentdb" // nolint: staticcheck
 	"github.com/hashicorp/go-azure-helpers/lang/pointer"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/cosmosdb"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/tf/pluginsdk"
 )
 
-func ExpandCosmosDbAutoscaleSettingsLegacy(d *pluginsdk.ResourceData) *documentdb.AutoscaleSettings {
-	i := d.Get("autoscale_settings").([]interface{})
-	if len(i) == 0 || i[0] == nil {
-		log.Printf("[DEBUG] Cosmos DB autoscale settings are not set on the resource")
-		return nil
-	}
-	input := i[0].(map[string]interface{})
-
-	autoscaleSettings := documentdb.AutoscaleSettings{}
-
-	if maxThroughput, ok := input["max_throughput"].(int); ok {
-		autoscaleSettings.MaxThroughput = pointer.To(int32(maxThroughput))
-	}
-
-	return &autoscaleSettings
-}
-
 func ExpandCosmosDbAutoscaleSettings(d *pluginsdk.ResourceData) *cosmosdb.AutoScaleSettings {
 	i := d.Get("autoscale_settings").([]interface{})
 	if len(i) == 0 || i[0] == nil {
@@ -46,34 +28,6 @@ func ExpandCosmosDbAutoscaleSettings(d *pluginsdk.ResourceData) *cosmosdb.AutoSc
 	return &autoscaleSettings
 }
 
-func FlattenCosmosDbAutoscaleSettingsLegacy(throughputResponse documentdb.ThroughputSettingsGetResults) []interface{} {
-	results := make([]interface{}, 0)
-
-	props := throughputResponse.ThroughputSettingsGetProperties
-	if props == nil {
-		return results
-	}
-
-	res := props.Resource
-	if res == nil {
-		return results
-	}
-
-	autoscaleSettings := res.AutoscaleSettings
-	if autoscaleSettings == nil {
-		log.Printf("[DEBUG] Cosmos DB autoscale settings are not set on the throughput response")
-		return results
-	}
-
-	result := make(map[string]interface{})
-
-	if autoscaleSettings.MaxThroughput != nil {
-		result["max_throughput"] = autoscaleSettings.MaxThroughput
-	}
-
-	return append(results, result)
-}
-
 func FlattenCosmosDbAutoscaleSettings(throughputResponse cosmosdb.ThroughputSettingsGetResults) []interface{} {
 	results := make([]interface{}, 0)
 
@@ -100,18 +54,6 @@ func FlattenCosmosDbAutoscaleSettings(throughputResponse cosmosdb.ThroughputSett
 	return append(results, result)
 }
 
-func ExpandCosmosDbAutoscaleSettingsResourceLegacy(d *pluginsdk.ResourceData) *documentdb.AutoscaleSettingsResource {
-	autoscaleSettings := ExpandCosmosDbAutoscaleSettingsLegacy(d)
-
-	if autoscaleSettings == nil {
-		return nil
-	}
-
-	return &documentdb.AutoscaleSettingsResource{
-		MaxThroughput: autoscaleSettings.MaxThroughput,
-	}
-}
-
 func ExpandCosmosDbAutoscaleSettingsResource(d *pluginsdk.ResourceData) *cosmosdb.AutoscaleSettingsResource {
 	autoscaleSettings := ExpandCosmosDbAutoscaleSettings(d)
 

internal/services/cosmos/common/schema.go

@@ -4,7 +4,7 @@
 package common
 
 import (
-	"github.com/Azure/azure-sdk-for-go/services/cosmos-db/mgmt/2021-10-15/documentdb" // nolint: staticcheck
+	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/cosmosdb"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/services/cosmos/validate"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/tf/pluginsdk"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/tf/suppress"
@@ -109,10 +109,10 @@ func CosmosDbIndexingPolicySchema() *pluginsdk.Schema {
 				"indexing_mode": {
 					Type:     pluginsdk.TypeString,
 					Optional: true,
-					Default:  documentdb.IndexingModeConsistent,
+					Default:  cosmosdb.IndexingModeConsistent,
 					ValidateFunc: validation.StringInSlice([]string{
-						string(documentdb.IndexingModeConsistent),
-						string(documentdb.IndexingModeNone),
+						string(cosmosdb.IndexingModeConsistent),
+						string(cosmosdb.IndexingModeNone),
 					}, false),
 				},
 
@@ -163,8 +163,8 @@ func ConflictResolutionPolicy() *pluginsdk.Schema {
 					Type:     pluginsdk.TypeString,
 					Required: true,
 					ValidateFunc: validation.StringInSlice([]string{
-						string(documentdb.ConflictResolutionModeLastWriterWins),
-						string(documentdb.ConflictResolutionModeCustom),
+						string(cosmosdb.ConflictResolutionModeLastWriterWins),
+						string(cosmosdb.ConflictResolutionModeCustom),
 					}, false),
 				},
 
@@ -208,8 +208,8 @@ func CosmosDbIndexingPolicyCompositeIndexSchema() *pluginsdk.Schema {
 								// Workaround for Azure/azure-rest-api-specs#11222
 								DiffSuppressFunc: suppress.CaseDifference,
 								ValidateFunc: validation.StringInSlice([]string{
-									string(documentdb.CompositePathSortOrderAscending),
-									string(documentdb.CompositePathSortOrderDescending),
+									string(cosmosdb.CompositePathSortOrderAscending),
+									string(cosmosdb.CompositePathSortOrderDescending),
 								}, true),
 							},
 						},

internal/services/cosmos/common/throughput.go

@@ -6,26 +6,11 @@ package common
 import (
 	"fmt"
 
-	"github.com/Azure/azure-sdk-for-go/services/cosmos-db/mgmt/2021-10-15/documentdb" // nolint: staticcheck
 	"github.com/hashicorp/go-azure-helpers/lang/pointer"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/cosmosdb"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/tf/pluginsdk"
 )
 
-func GetThroughputFromResultLegacy(throughputResponse documentdb.ThroughputSettingsGetResults) *int32 {
-	props := throughputResponse.ThroughputSettingsGetProperties
-	if props == nil {
-		return nil
-	}
-
-	res := props.Resource
-	if res == nil {
-		return nil
-	}
-
-	return res.Throughput
-}
-
 func GetThroughputFromResult(throughputResponse cosmosdb.ThroughputSettingsGetResults) *int64 {
 	props := throughputResponse.Properties
 	if props == nil {
@@ -40,35 +25,11 @@ func GetThroughputFromResult(throughputResponse cosmosdb.ThroughputSettingsGetRe
 	return res.Throughput
 }
 
-func ConvertThroughputFromResourceDataLegacy(throughput interface{}) *int32 {
-	return pointer.To(int32(throughput.(int)))
-}
-
 func ConvertThroughputFromResourceData(throughput interface{}) *int64 {
 	return pointer.To(int64(throughput.(int)))
 }
 
-func ExpandCosmosDBThroughputSettingsUpdateParametersLegacy(d *pluginsdk.ResourceData) *documentdb.ThroughputSettingsUpdateParameters {
-	throughputParameters := documentdb.ThroughputSettingsUpdateParameters{
-		ThroughputSettingsUpdateProperties: &documentdb.ThroughputSettingsUpdateProperties{
-			Resource: &documentdb.ThroughputSettingsResource{},
-		},
-	}
-
-	if v, exists := d.GetOk("throughput"); exists {
-		throughputParameters.Resource.Throughput = ConvertThroughputFromResourceDataLegacy(v)
-	}
-
-	if _, hasAutoscaleSettings := d.GetOk("autoscale_settings"); hasAutoscaleSettings {
-		// If updating the autoscale throughput, set the manual throughput to nil to ensure the autoscale throughput is applied
-		throughputParameters.Resource.Throughput = nil
-		throughputParameters.Resource.AutoscaleSettings = ExpandCosmosDbAutoscaleSettingsResourceLegacy(d)
-	}
-
-	return &throughputParameters
-}
-
-func ExpandCosmosDBThroughputSettingsUpdateParameters(d *pluginsdk.ResourceData) *cosmosdb.ThroughputSettingsUpdateParameters {
+func ExpandCosmosDBThroughputSettingsUpdateParameters(d *pluginsdk.ResourceData) cosmosdb.ThroughputSettingsUpdateParameters {
 	throughputParameters := cosmosdb.ThroughputSettingsUpdateParameters{
 		Properties: cosmosdb.ThroughputSettingsUpdateProperties{
 			Resource: cosmosdb.ThroughputSettingsResource{},
@@ -85,14 +46,7 @@ func ExpandCosmosDBThroughputSettingsUpdateParameters(d *pluginsdk.ResourceData)
 		throughputParameters.Properties.Resource.AutoScaleSettings = ExpandCosmosDbAutoscaleSettingsResource(d)
 	}
 
-	return &throughputParameters
-}
-
-func SetResourceDataThroughputFromResponseLegacy(throughputResponse documentdb.ThroughputSettingsGetResults, d *pluginsdk.ResourceData) {
-	d.Set("throughput", GetThroughputFromResultLegacy(throughputResponse))
-
-	autoscaleSettings := FlattenCosmosDbAutoscaleSettingsLegacy(throughputResponse)
-	d.Set("autoscale_settings", autoscaleSettings)
+	return throughputParameters
 }
 
 func SetResourceDataThroughputFromResponse(throughputResponse cosmosdb.ThroughputSettingsGetResults, d *pluginsdk.ResourceData) {

internal/services/cosmos/cosmosdb.go

@@ -1,13 +1,18 @@
 package cosmos
 
-import "github.com/Azure/azure-sdk-for-go/services/cosmos-db/mgmt/2021-10-15/documentdb" // nolint: staticcheck
-
-func isServerlessCapacityMode(accResp documentdb.DatabaseAccountGetResults) bool {
-	if props := accResp.DatabaseAccountGetProperties; props != nil && props.Capabilities != nil {
-		for _, v := range *props.Capabilities {
-			if v.Name != nil && *v.Name == "EnableServerless" {
-				return true
-			}
+import (
+	"github.com/hashicorp/go-azure-helpers/lang/pointer"
+	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/cosmosdb"
+)
+
+func isServerlessCapacityMode(input *cosmosdb.DatabaseAccountGetResults) bool {
+	if input == nil || input.Properties == nil || input.Properties.Capabilities == nil {
+		return false
+	}
+
+	for _, v := range *input.Properties.Capabilities {
+		if pointer.From(v.Name) == "EnableServerless" {
+			return true
 		}
 	}
 

internal/services/cosmos/cosmosdb_account_resource.go

@@ -22,6 +22,7 @@ import (
 	"github.com/hashicorp/go-azure-helpers/resourcemanager/location"
 	"github.com/hashicorp/go-azure-helpers/resourcemanager/tags"
 	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/cosmosdb"
+	"github.com/hashicorp/go-azure-sdk/resource-manager/cosmosdb/2024-08-15/restorables"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-provider-azurerm/helpers/azure"
 	"github.com/hashicorp/terraform-provider-azurerm/helpers/tf"
@@ -30,7 +31,6 @@ import (
 	"github.com/hashicorp/terraform-provider-azurerm/internal/sdk"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/services/cosmos/common"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/services/cosmos/migration"
-	"github.com/hashicorp/terraform-provider-azurerm/internal/services/cosmos/parse"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/services/cosmos/validate"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/tf/pluginsdk"
 	"github.com/hashicorp/terraform-provider-azurerm/internal/tf/validation"
@@ -180,7 +180,7 @@ func resourceCosmosDbAccount() *pluginsdk.Resource {
 		),
 
 		Importer: pluginsdk.ImporterValidatingResourceId(func(id string) error {
-			_, err := parse.DatabaseAccountID(id)
+			_, err := cosmosdb.ParseDatabaseAccountID(id)
 			return err
 		}),
 
@@ -591,7 +591,7 @@ func resourceCosmosDbAccount() *pluginsdk.Resource {
 							Type:         pluginsdk.TypeString,
 							Required:     true,
 							ForceNew:     true,
-							ValidateFunc: validate.RestorableDatabaseAccountID,
+							ValidateFunc: restorables.ValidateRestorableDatabaseAccountID,
 						},
 
 						"restore_timestamp_in_utc": {