Support autoscaling Cosmos DB databases and containers

azurerm/helpers/validate/cosmos.go

@@ -42,3 +42,28 @@ func CosmosThroughput(v interface{}, k string) (warnings []string, errors []erro
 
 	return warnings, errors
 }
+
+func CosmosMaxThroughput(i interface{}, k string) (warnings []string, errors []error) {
+	v, ok := i.(int)
+	if !ok {
+		errors = append(errors, fmt.Errorf("expected type of %q to be int", k))
+		return
+	}
+
+	if v < 4000 {
+		errors = append(errors, fmt.Errorf(
+			"%s must be a minimum of 4000", k))
+	}
+
+	if v > 1000000 {
+		errors = append(errors, fmt.Errorf(
+			"%s must be a maximum of 1000000", k))
+	}
+
+	if v%1000 != 0 {
+		errors = append(errors, fmt.Errorf(
+			"%q must be set in increments of 1000", k))
+	}
+
+	return warnings, errors
+}

azurerm/helpers/validate/cosmos_test.go

@@ -98,3 +98,54 @@ func TestCosmosThroughput(t *testing.T) {
 		}
 	}
 }
+
+func TestCosmosMaxThroughput(t *testing.T) {
+	cases := []struct {
+		Value  interface{}
+		Errors int
+	}{
+		{
+			Value:  400,
+			Errors: 2,
+		},
+		{
+			Value:  1000,
+			Errors: 1,
+		},
+		{
+			Value:  4000,
+			Errors: 0,
+		},
+		{
+			Value:  4001,
+			Errors: 1,
+		},
+		{
+			Value:  10000,
+			Errors: 0,
+		},
+		{
+			Value:  54000,
+			Errors: 0,
+		},
+		{
+			Value:  1000000,
+			Errors: 0,
+		},
+		{
+			Value:  1100000,
+			Errors: 1,
+		},
+		{
+			Value:  "400",
+			Errors: 1,
+		},
+	}
+
+	for _, tc := range cases {
+		_, errors := CosmosMaxThroughput(tc.Value, "throughput")
+		if len(errors) != tc.Errors {
+			t.Fatalf("Expected CosmosMaxThroughput to trigger '%d' errors for '%d' - got '%d'", tc.Errors, tc.Value, len(errors))
+		}
+	}
+}

azurerm/internal/services/cosmos/common/autoscale_settings.go

@@ -0,0 +1,62 @@
+package common
+
+import (
+	"log"
+
+	"github.com/Azure/azure-sdk-for-go/services/cosmos-db/mgmt/2020-04-01/documentdb"
+	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
+	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/utils"
+)
+
+func ExpandCosmosDbAutoscaleSettings(d *schema.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 = utils.Int32(int32(maxThroughput))
+	}
+
+	return &autoscaleSettings
+}
+
+func FlattenCosmosDbAutoscaleSettings(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 ExpandCosmosDbAutoscaleSettingsResource(d *schema.ResourceData) *documentdb.AutoscaleSettingsResource {
+	autoscaleSettings := ExpandCosmosDbAutoscaleSettings(d)
+	autoscaleSettingResource := documentdb.AutoscaleSettingsResource{}
+
+	autoscaleSettingResource.MaxThroughput = autoscaleSettings.MaxThroughput
+	return &autoscaleSettingResource
+}

azurerm/internal/services/cosmos/common/schema.go

@@ -0,0 +1,39 @@
+package common
+
+import (
+	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
+	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/helpers/validate"
+)
+
+func DatabaseAutoscaleSettingsSchema() *schema.Schema {
+	return &schema.Schema{
+		Type:     schema.TypeList,
+		Optional: true,
+		MaxItems: 1,
+		Elem: &schema.Resource{
+			Schema: map[string]*schema.Schema{
+				"max_throughput": {
+					Type:          schema.TypeInt,
+					Optional:      true,
+					Computed:      true,
+					ConflictsWith: []string{"throughput"},
+					ValidateFunc:  validate.CosmosMaxThroughput,
+				},
+			},
+		},
+	}
+}
+
+func ContainerAutoscaleSettingsSchema() *schema.Schema {
+	autoscaleSettingsDatabaseSchema := DatabaseAutoscaleSettingsSchema()
+	autoscaleSettingsDatabaseSchema.RequiredWith = []string{"partition_key_path"}
+
+	return autoscaleSettingsDatabaseSchema
+}
+
+func MongoCollectionAutoscaleSettingsSchema() *schema.Schema {
+	autoscaleSettingsDatabaseSchema := DatabaseAutoscaleSettingsSchema()
+	autoscaleSettingsDatabaseSchema.RequiredWith = []string{"shard_key"}
+
+	return autoscaleSettingsDatabaseSchema
+}

azurerm/internal/services/cosmos/common/throughput.go

@@ -1,7 +1,10 @@
 package common
 
 import (
+	"fmt"
+
 	"github.com/Azure/azure-sdk-for-go/services/cosmos-db/mgmt/2020-04-01/documentdb"
+	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/utils"
 )
 
@@ -22,3 +25,42 @@ func GetThroughputFromResult(throughputResponse documentdb.ThroughputSettingsGet
 func ConvertThroughputFromResourceData(throughput interface{}) *int32 {
 	return utils.Int32(int32(throughput.(int)))
 }
+
+func ExpandCosmosDBThroughputSettingsUpdateParameters(d *schema.ResourceData) *documentdb.ThroughputSettingsUpdateParameters {
+	throughputParameters := documentdb.ThroughputSettingsUpdateParameters{
+		ThroughputSettingsUpdateProperties: &documentdb.ThroughputSettingsUpdateProperties{
+			Resource: &documentdb.ThroughputSettingsResource{},
+		},
+	}
+
+	if v, exists := d.GetOk("throughput"); exists {
+		throughputParameters.ThroughputSettingsUpdateProperties.Resource.Throughput = ConvertThroughputFromResourceData(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.ThroughputSettingsUpdateProperties.Resource.Throughput = nil
+		throughputParameters.ThroughputSettingsUpdateProperties.Resource.AutoscaleSettings = ExpandCosmosDbAutoscaleSettingsResource(d)
+	}
+
+	return &throughputParameters
+}
+
+func SetResourceDataThroughputFromResponse(throughputResponse documentdb.ThroughputSettingsGetResults, d *schema.ResourceData) {
+	d.Set("throughput", GetThroughputFromResult(throughputResponse))
+
+	autoscaleSettings := FlattenCosmosDbAutoscaleSettings(throughputResponse)
+	d.Set("autoscale_settings", autoscaleSettings)
+}
+
+func CheckForChangeFromAutoscaleAndManualThroughput(d *schema.ResourceData) error {
+	if d.HasChange("throughput") && d.HasChange("autoscale_settings") {
+		return fmt.Errorf("switching between autoscale and manually provisioned throughput via Terraform is not supported at this time")
+	}
+
+	return nil
+}
+
+func HasThroughputChange(d *schema.ResourceData) bool {
+	return d.HasChanges("throughput", "autoscale_settings")
+}

azurerm/internal/services/cosmos/cosmosdb_cassandra_keyspace_resource.go

@@ -69,6 +69,8 @@ func resourceArmCosmosDbCassandraKeyspace() *schema.Resource {
 				Computed:     true,
 				ValidateFunc: validate.CosmosThroughput,
 			},
+
+			"autoscale_settings": common.DatabaseAutoscaleSettingsSchema(),
 		},
 	}
 }
@@ -105,7 +107,13 @@ func resourceArmCosmosDbCassandraKeyspaceCreate(d *schema.ResourceData, meta int
 	}
 
 	if throughput, hasThroughput := d.GetOk("throughput"); hasThroughput {
-		db.CassandraKeyspaceCreateUpdateProperties.Options.Throughput = common.ConvertThroughputFromResourceData(throughput)
+		if throughput != 0 {
+			db.CassandraKeyspaceCreateUpdateProperties.Options.Throughput = common.ConvertThroughputFromResourceData(throughput)
+		}
+	}
+
+	if _, hasAutoscaleSettings := d.GetOk("autoscale_settings"); hasAutoscaleSettings {
+		db.CassandraKeyspaceCreateUpdateProperties.Options.AutoscaleSettings = common.ExpandCosmosDbAutoscaleSettings(d)
 	}
 
 	future, err := client.CreateUpdateCassandraKeyspace(ctx, resourceGroup, account, name, db)
@@ -141,6 +149,11 @@ func resourceArmCosmosDbCassandraKeyspaceUpdate(d *schema.ResourceData, meta int
 		return err
 	}
 
+	err = common.CheckForChangeFromAutoscaleAndManualThroughput(d)
+	if err != nil {
+		return fmt.Errorf("Error updating Cosmos Cassandra Keyspace %q (Account: %q) - %+v", id.Name, id.Account, err)
+	}
+
 	db := documentdb.CassandraKeyspaceCreateUpdateParameters{
 		CassandraKeyspaceCreateUpdateProperties: &documentdb.CassandraKeyspaceCreateUpdateProperties{
 			Resource: &documentdb.CassandraKeyspaceResource{
@@ -159,16 +172,9 @@ func resourceArmCosmosDbCassandraKeyspaceUpdate(d *schema.ResourceData, meta int
 		return fmt.Errorf("Error waiting on create/update future for Cosmos Cassandra Keyspace %q (Account: %q): %+v", id.ResourceGroup, id.Account, err)
 	}
 
-	if d.HasChange("throughput") {
-		throughputParameters := documentdb.ThroughputSettingsUpdateParameters{
-			ThroughputSettingsUpdateProperties: &documentdb.ThroughputSettingsUpdateProperties{
-				Resource: &documentdb.ThroughputSettingsResource{
-					Throughput: common.ConvertThroughputFromResourceData(d.Get("throughput")),
-				},
-			},
-		}
-
-		throughputFuture, err := client.UpdateCassandraKeyspaceThroughput(ctx, id.ResourceGroup, id.Account, id.Name, throughputParameters)
+	if common.HasThroughputChange(d) {
+		throughputParameters := common.ExpandCosmosDBThroughputSettingsUpdateParameters(d)
+		throughputFuture, err := client.UpdateCassandraKeyspaceThroughput(ctx, id.ResourceGroup, id.Account, id.Name, *throughputParameters)
 		if err != nil {
 			if response.WasNotFound(throughputFuture.Response()) {
 				return fmt.Errorf("Error setting Throughput for Cosmos Cassandra Keyspace %q (Account: %q): %+v - "+
@@ -219,9 +225,10 @@ func resourceArmCosmosDbCassandraKeyspaceRead(d *schema.ResourceData, meta inter
 			return fmt.Errorf("Error reading Throughput on Cosmos Cassandra Keyspace %q (Account: %q): %+v", id.Name, id.Account, err)
 		} else {
 			d.Set("throughput", nil)
+			d.Set("autoscale_settings", nil)
 		}
 	} else {
-		d.Set("throughput", common.GetThroughputFromResult(throughputResp))
+		common.SetResourceDataThroughputFromResponse(throughputResp, d)
 	}
 
 	return nil

azurerm/internal/services/cosmos/cosmosdb_gremlin_database_resource.go

@@ -69,6 +69,8 @@ func resourceArmCosmosGremlinDatabase() *schema.Resource {
 				Computed:     true,
 				ValidateFunc: validate.CosmosThroughput,
 			},
+
+			"autoscale_settings": common.DatabaseAutoscaleSettingsSchema(),
 		},
 	}
 }
@@ -105,7 +107,13 @@ func resourceArmCosmosGremlinDatabaseCreate(d *schema.ResourceData, meta interfa
 	}
 
 	if throughput, hasThroughput := d.GetOk("throughput"); hasThroughput {
-		db.GremlinDatabaseCreateUpdateProperties.Options.Throughput = common.ConvertThroughputFromResourceData(throughput)
+		if throughput != 0 {
+			db.GremlinDatabaseCreateUpdateProperties.Options.Throughput = common.ConvertThroughputFromResourceData(throughput)
+		}
+	}
+
+	if _, hasAutoscaleSettings := d.GetOk("autoscale_settings"); hasAutoscaleSettings {
+		db.GremlinDatabaseCreateUpdateProperties.Options.AutoscaleSettings = common.ExpandCosmosDbAutoscaleSettings(d)
 	}
 
 	future, err := client.CreateUpdateGremlinDatabase(ctx, resourceGroup, account, name, db)
@@ -141,6 +149,11 @@ func resourceArmCosmosGremlinDatabaseUpdate(d *schema.ResourceData, meta interfa
 		return err
 	}
 
+	err = common.CheckForChangeFromAutoscaleAndManualThroughput(d)
+	if err != nil {
+		return fmt.Errorf("Error updating Cosmos Gremlin Database %q (Account: %q) - %+v", id.Name, id.Account, err)
+	}
+
 	db := documentdb.GremlinDatabaseCreateUpdateParameters{
 		GremlinDatabaseCreateUpdateProperties: &documentdb.GremlinDatabaseCreateUpdateProperties{
 			Resource: &documentdb.GremlinDatabaseResource{
@@ -159,16 +172,9 @@ func resourceArmCosmosGremlinDatabaseUpdate(d *schema.ResourceData, meta interfa
 		return fmt.Errorf("Error waiting on create/update future for Cosmos Gremlin Database %q (Account: %q): %+v", id.Name, id.Account, err)
 	}
 
-	if d.HasChange("throughput") {
-		throughputParameters := documentdb.ThroughputSettingsUpdateParameters{
-			ThroughputSettingsUpdateProperties: &documentdb.ThroughputSettingsUpdateProperties{
-				Resource: &documentdb.ThroughputSettingsResource{
-					Throughput: common.ConvertThroughputFromResourceData(d.Get("throughput")),
-				},
-			},
-		}
-
-		throughputFuture, err := client.UpdateGremlinDatabaseThroughput(ctx, id.ResourceGroup, id.Account, id.Name, throughputParameters)
+	if common.HasThroughputChange(d) {
+		throughputParameters := common.ExpandCosmosDBThroughputSettingsUpdateParameters(d)
+		throughputFuture, err := client.UpdateGremlinDatabaseThroughput(ctx, id.ResourceGroup, id.Account, id.Name, *throughputParameters)
 		if err != nil {
 			if response.WasNotFound(throughputFuture.Response()) {
 				return fmt.Errorf("Error setting Throughput for Cosmos Gremlin Database %q (Account: %q): %+v - "+
@@ -223,9 +229,10 @@ func resourceArmCosmosGremlinDatabaseRead(d *schema.ResourceData, meta interface
 			return fmt.Errorf("Error reading Throughput on Cosmos Gremlin Database %q (Account: %q): %+v", id.Name, id.Account, err)
 		} else {
 			d.Set("throughput", nil)
+			d.Set("autoscale_settings", nil)
 		}
 	} else {
-		d.Set("throughput", common.GetThroughputFromResult(throughputResp))
+		common.SetResourceDataThroughputFromResponse(throughputResp, d)
 	}
 
 	return nil