allow to push additional dimension specified by command-line option (#1) (#4)

* allow to push additional dimension specified by command-line option (#1)

* fix poor english

* update README.md

README.md

@@ -21,6 +21,7 @@ Command-line arguments take precedence over ENV vars
 | cert_path                    | CERT_PATH                    | Path to SSL Certificate file (when using SSL for `prometheus_scrape_url`)     |
 | keyPath                      | KEY_PATH                     | Path to Key file (when using SSL for `prometheus_scrape_url`)                 |
 | accept_invalid_cert          | ACCEPT_INVALID_CERT          | Accept any certificate during TLS handshake. Insecure, use only for testing   |
+| additional_dimension         | ADDITIONAL_DIMENSION         | Additional dimension specified by NAME=VALUE                                  |
 
 
 __NOTE__: If AWS credentials are not provided in the command-line arguments (`aws_access_key_id` and `aws_secret_access_key`)

main.go

@@ -7,6 +7,7 @@ import (
 	"log"
 	"os"
 	"strconv"
+	"strings"
 	"time"
 )
 
@@ -21,6 +22,7 @@ var (
 	certPath                 = flag.String("cert_path", os.Getenv("CERT_PATH"), "Path to SSL Certificate file (when using SSL for `prometheus_scrape_url`)")
 	keyPath                  = flag.String("key_path", os.Getenv("KEY_PATH"), "Path to Key file (when using SSL for `prometheus_scrape_url`)")
 	skipServerCertCheck      = flag.String("accept_invalid_cert", os.Getenv("ACCEPT_INVALID_CERT"), "Accept any certificate during TLS handshake. Insecure, use only for testing")
+	additionalDimension      = flag.String("additional_dimension", os.Getenv("ADDITIONAL_DIMENSION"), "Additional dimension specified by NAME=VALUE")
 )
 
 func main() {
@@ -52,6 +54,15 @@ func main() {
 		}
 	}
 
+	var additionalDimensions = map[string]string{}
+	if *additionalDimension != "" {
+		kv := strings.SplitN(*additionalDimension, "=", 2)
+		if len(kv) != 2 {
+			log.Fatal("prometheus-to-cloudwatch: Error: -additionalDimension must be formated as NAME=VALUE")
+		}
+		additionalDimensions[kv[0]] = kv[1]
+	}
+
 	config := &Config{
 		CloudWatchNamespace:           *cloudWatchNamespace,
 		CloudWatchRegion:              *cloudWatchRegion,
@@ -61,6 +72,7 @@ func main() {
 		PrometheusSkipServerCertCheck: skipCertCheck,
 		AwsAccessKeyId:                *awsAccessKeyId,
 		AwsSecretAccessKey:            *awsSecretAccessKey,
+		AdditionalDimensions:          additionalDimensions,
 	}
 
 	if *prometheusScrapeInterval != "" {

prometheus_to_cloudwatch.go

@@ -59,6 +59,9 @@ type Config struct {
 
 	// Accept any certificate during TLS handshake. Insecure, use only for testing
 	PrometheusSkipServerCertCheck bool
+
+	// Additional dimensions to send to CloudWatch
+	AdditionalDimensions map[string]string
 }
 
 // Bridge pushes metrics to AWS CloudWatch
@@ -70,6 +73,7 @@ type Bridge struct {
 	prometheusCertPath            string
 	prometheusKeyPath             string
 	prometheusSkipServerCertCheck bool
+	additionalDimensions          map[string]string
 }
 
 // NewBridge initializes and returns a pointer to a Bridge using the
@@ -90,6 +94,7 @@ func NewBridge(c *Config) (*Bridge, error) {
 	b.prometheusCertPath = c.PrometheusCertPath
 	b.prometheusKeyPath = c.PrometheusKeyPath
 	b.prometheusSkipServerCertCheck = c.PrometheusSkipServerCertCheck
+	b.additionalDimensions = c.AdditionalDimensions
 
 	if c.CloudWatchPublishInterval > 0 {
 		b.cloudWatchPublishInterval = c.CloudWatchPublishInterval
@@ -174,7 +179,7 @@ func (b *Bridge) publishMetricsToCloudWatch(mfs []*dto.MetricFamily) error {
 
 	for _, s := range vec {
 		name := getName(s.Metric)
-		data = appendDatum(data, name, s)
+		data = appendDatum(data, name, s, b)
 
 		if len(data) == batchSize {
 			if err := b.flush(data); err != nil {
@@ -199,7 +204,7 @@ func (b *Bridge) flush(data []*cloudwatch.MetricDatum) error {
 	return nil
 }
 
-func appendDatum(data []*cloudwatch.MetricDatum, name string, s *model.Sample) []*cloudwatch.MetricDatum {
+func appendDatum(data []*cloudwatch.MetricDatum, name string, s *model.Sample, b *Bridge) []*cloudwatch.MetricDatum {
 	metric := s.Metric
 
 	if len(metric) == 0 {
@@ -211,7 +216,7 @@ func appendDatum(data []*cloudwatch.MetricDatum, name string, s *model.Sample) [
 	d.SetMetricName(name).
 		SetValue(float64(s.Value)).
 		SetTimestamp(s.Timestamp.Time()).
-		SetDimensions(getDimensions(metric)).
+		SetDimensions(append(getDimensions(metric, 10-len(b.additionalDimensions)), getAdditionalDimensions(b)...)).
 		SetStorageResolution(getResolution(metric)).
 		SetUnit(getUnit(metric))
 
@@ -227,7 +232,7 @@ func getName(m model.Metric) string {
 
 // getDimensions returns up to 10 dimensions for the provided metric - one for each label (except the __name__ label)
 // If a metric has more than 10 labels, it attempts to behave deterministically and returning the first 10 labels as dimensions
-func getDimensions(m model.Metric) []*cloudwatch.Dimension {
+func getDimensions(m model.Metric, num int) []*cloudwatch.Dimension {
 	if len(m) == 0 {
 		return make([]*cloudwatch.Dimension, 0)
 	} else if _, ok := m[model.MetricNameLabel]; len(m) == 1 && ok {
@@ -253,13 +258,21 @@ func getDimensions(m model.Metric) []*cloudwatch.Dimension {
 		}
 	}
 
-	if len(dims) > 10 {
-		dims = dims[:10]
+	if len(dims) > num {
+		dims = dims[:num]
 	}
 
 	return dims
 }
 
+func getAdditionalDimensions(b *Bridge) []*cloudwatch.Dimension {
+	dims := make([]*cloudwatch.Dimension, 0, len(b.additionalDimensions))
+	for k, v := range b.additionalDimensions {
+		dims = append(dims, new(cloudwatch.Dimension).SetName(k).SetValue(v))
+	}
+	return dims
+}
+
 // Returns 1 if the metric contains a __cw_high_res label, otherwise returns 60
 func getResolution(m model.Metric) int64 {
 	if _, ok := m[cwHighResLabel]; ok {