CosmosDB: Add throughput configuration (#6693)

* CosmosDB: Add throughput configuration * CR Fixes * CR Fixes
treeverse · Oct 4, 2023 · 6e4a289 · 6e4a289
1 parent e901590
commit 6e4a289
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 10 deletions.
diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md
@@ -65,6 +65,8 @@ This reference uses `.` to denote the nesting of values.
     + `database.cosmosdb.endpoint` `(string : "")` - CosmosDB account endpoint, e.g. `https://<account>.documents.azure.com/`.
     + `database.cosmosdb.database` `(string : "")` - CosmosDB database name.
     + `database.cosmosdb.container` `(string : "")` - CosmosDB container name.
+    + `database.cosmosdb.throughput` `(int32 : )` - CosmosDB container's RU/s. If not set - the default CosmosDB container throughput is used. 
+    + `database.cosmosdb.autoscale` `(bool : false)` - If set, CosmosDB container throughput is autoscaled (See CosmosDB docs for minimum throughput requirement). Otherwise, uses "Manual" mode ([Docs](https://learn.microsoft.com/en-us/azure/cosmos-db/provision-throughput-autoscale)).
   + `database.local` - Configuration section when using `database.type="local"`
     + `database.local.path` `(string : "~/lakefs/metadata")` - Local path on the filesystem to store embedded KV metadata, like branches and uncommitted entries
     + `database.local.sync_writes` `(bool: true)` - Ensure each write is written to the disk. Disable to increase performance

diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -172,10 +172,12 @@ type Config struct {
 		} `mapstructure:"dynamodb"`
 
 		CosmosDB *struct {
-			Key       SecureString `mapstructure:"key"`
-			Endpoint  string       `mapstructure:"endpoint"`
-			Database  string       `mapstructure:"database"`
-			Container string       `mapstructure:"container"`
+			Key        SecureString `mapstructure:"key"`
+			Endpoint   string       `mapstructure:"endpoint"`
+			Database   string       `mapstructure:"database"`
+			Container  string       `mapstructure:"container"`
+			Throughput int32        `mapstructure:"throughput"`
+			Autoscale  bool         `mapstructure:"autoscale"`
 		} `mapstructure:"cosmosdb"`
 	}
 

diff --git a/pkg/kv/cosmosdb/store.go b/pkg/kv/cosmosdb/store.go
@@ -129,6 +129,17 @@ func getOrCreateDatabase(ctx context.Context, client *azcosmos.Client, params *k
 }
 
 func getOrCreateContainer(ctx context.Context, dbClient *azcosmos.DatabaseClient, params *kvparams.CosmosDB) (*azcosmos.ContainerClient, error) {
+	var opts *azcosmos.CreateContainerOptions
+	if params.Throughput > 0 {
+		var throughputProperties azcosmos.ThroughputProperties
+		if params.Autoscale {
+			throughputProperties = azcosmos.NewAutoscaleThroughputProperties(params.Throughput)
+		} else {
+			throughputProperties = azcosmos.NewManualThroughputProperties(params.Throughput)
+		}
+		opts = &azcosmos.CreateContainerOptions{ThroughputProperties: &throughputProperties}
+	}
+
 	_, err := dbClient.CreateContainer(ctx,
 		azcosmos.ContainerProperties{
 			ID: params.Container,
@@ -144,7 +155,7 @@ func getOrCreateContainer(ctx context.Context, dbClient *azcosmos.DatabaseClient
 				IncludedPaths: []azcosmos.IncludedPath{{Path: "/*"}},
 				ExcludedPaths: []azcosmos.ExcludedPath{{Path: "/value/?"}},
 			},
-		}, nil)
+		}, opts)
 	if err != nil {
 		if errStatusCode(err) != http.StatusConflict {
 			return nil, fmt.Errorf("creating container: %w", err)

diff --git a/pkg/kv/kvparams/database.go b/pkg/kv/kvparams/database.go
@@ -63,11 +63,12 @@ type DynamoDB struct {
 }
 
 type CosmosDB struct {
-	Key       string
-	Endpoint  string
-	Database  string
-	Container string
-
+	Key        string
+	Endpoint   string
+	Database   string
+	Container  string
+	Throughput int32
+	Autoscale  bool
 	// These values should only be set to false for testing purposes using the CosmosDB emulator
 	Client            *http.Client
 	StrongConsistency bool
@@ -111,11 +112,16 @@ func NewConfig(cfg *config.Config) (Config, error) {
 	}
 
 	if cfg.Database.CosmosDB != nil {
+		if cfg.Database.CosmosDB.Autoscale && cfg.Database.CosmosDB.Throughput == 0 {
+			return Config{}, fmt.Errorf("enabling autoscale requires setting the throughput param: %w", config.ErrBadConfiguration)
+		}
 		p.CosmosDB = &CosmosDB{
 			Key:               cfg.Database.CosmosDB.Key.SecureValue(),
 			Endpoint:          cfg.Database.CosmosDB.Endpoint,
 			Database:          cfg.Database.CosmosDB.Database,
 			Container:         cfg.Database.CosmosDB.Container,
+			Throughput:        cfg.Database.CosmosDB.Throughput,
+			Autoscale:         cfg.Database.CosmosDB.Autoscale,
 			StrongConsistency: true,
 		}
 	}