Only enable auto scale if passed ApplicationAutoScale url

Author: aaron7

cmd/table-manager/main.go

@@ -28,6 +28,10 @@ func main() {
 	util.RegisterFlags(&serverConfig, &storageConfig, &schemaConfig)
 	flag.Parse()
 
+	if schemaConfig.ChunkTables.WriteScaleEnabled && storageConfig.ApplicationAutoScaling.URL != nil {
+		log.Fatal("WriteScale is enabled but no ApplicationAutoScaling URL has been provided")
+	}
+
 	tableClient, err := storage.NewTableClient(storageConfig)
 	if err != nil {
 		log.Fatalf("Error initializing DynamoDB table client: %v", err)

pkg/chunk/aws_storage_client.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
+	"github.com/aws/aws-sdk-go/aws/client"
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/aws/session"
@@ -103,15 +104,17 @@ func init() {
 
 // DynamoDBConfig specifies config for a DynamoDB database.
 type DynamoDBConfig struct {
-	DynamoDB util.URLValue
-	APILimit float64
+	DynamoDB               util.URLValue
+	APILimit               float64
+	ApplicationAutoScaling util.URLValue
 }
 
 // RegisterFlags adds the flags required to config this to the given FlagSet
 func (cfg *DynamoDBConfig) RegisterFlags(f *flag.FlagSet) {
 	f.Var(&cfg.DynamoDB, "dynamodb.url", "DynamoDB endpoint URL with escaped Key and Secret encoded. "+
 		"If only region is specified as a host, proper endpoint will be deduced. Use inmemory:///<table-name> to use a mock in-memory implementation.")
 	f.Float64Var(&cfg.APILimit, "dynamodb.api-limit", 2.0, "DynamoDB table management requests per second limit.")
+	f.Var(&cfg.ApplicationAutoScaling, "applicationautoscaling.url", "ApplicationAutoscaling endpoint URL with escaped Key and Secret encoded.")
 }
 
 // AWSStorageConfig specifies config for storing data on AWS.
@@ -806,6 +809,15 @@ func recordDynamoError(tableName string, err error, operation string) {
 
 // dynamoClientFromURL creates a new DynamoDB client from a URL.
 func dynamoClientFromURL(awsURL *url.URL) (dynamodbiface.DynamoDBAPI, error) {
+	dynamoDBSession, err := awsSessionFromURL(awsURL)
+	if err != nil {
+		return nil, err
+	}
+	return dynamodb.New(dynamoDBSession), nil
+}
+
+// awsSessionFromURL creates a new aws session from a URL.
+func awsSessionFromURL(awsURL *url.URL) (client.ConfigProvider, error) {
 	if awsURL == nil {
 		return nil, fmt.Errorf("no URL specified for DynamoDB")
 	}
@@ -817,7 +829,7 @@ func dynamoClientFromURL(awsURL *url.URL) (dynamodbiface.DynamoDBAPI, error) {
 	if err != nil {
 		return nil, err
 	}
-	return dynamodb.New(session.New(config)), nil
+	return session.New(config), nil
 }
 
 // awsConfigFromURL returns AWS config from given URL. It expects escaped AWS Access key ID & Secret Access Key to be

pkg/chunk/dynamodb_table_client.go

@@ -2,19 +2,15 @@ package chunk
 
 import (
 	"fmt"
-	"net/url"
-	"strings"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
-	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/applicationautoscaling"
 	"github.com/aws/aws-sdk-go/service/applicationautoscaling/applicationautoscalingiface"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
 	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/common/log"
 	"golang.org/x/net/context"
 	"golang.org/x/time/rate"
 
@@ -49,9 +45,13 @@ func NewDynamoDBTableClient(cfg DynamoDBConfig) (TableClient, error) {
 		return nil, err
 	}
 
-	applicationAutoScaling, err := applicationAutoScalingClientFromURL(cfg.DynamoDB.URL)
-	if err != nil {
-		return nil, err
+	var applicationAutoScaling applicationautoscalingiface.ApplicationAutoScalingAPI
+	if cfg.ApplicationAutoScaling.URL != nil {
+		session, err := awsSessionFromURL(cfg.ApplicationAutoScaling.URL)
+		if err != nil {
+			return nil, err
+		}
+		applicationAutoScaling = applicationautoscaling.New(session)
 	}
 
 	return dynamoTableClient{
@@ -61,22 +61,6 @@ func NewDynamoDBTableClient(cfg DynamoDBConfig) (TableClient, error) {
 	}, nil
 }
 
-// applicationAutoScalingClientFromURL creates a new ApplicationAuthScaling client from a URL.
-func applicationAutoScalingClientFromURL(awsURL *url.URL) (applicationautoscalingiface.ApplicationAutoScalingAPI, error) {
-	if awsURL == nil {
-		return nil, fmt.Errorf("no URL specified for DynamoDB")
-	}
-	path := strings.TrimPrefix(awsURL.Path, "/")
-	if len(path) > 0 {
-		log.Warnf("Ignoring DynamoDB URL path: %v.", path)
-	}
-	config, err := awsConfigFromURL(awsURL)
-	if err != nil {
-		return nil, err
-	}
-	return applicationautoscaling.New(session.New(config)), nil
-}
-
 func (d dynamoTableClient) backoffAndRetry(ctx context.Context, fn func(context.Context) error) error {
 	if d.limiter != nil { // Tests will have a nil limiter.
 		d.limiter.Wait(ctx)