From 11e277d63405b1280768029a444ead62371e2730 Mon Sep 17 00:00:00 2001
From: Colin Liddle <48459220+colinliddlemyob@users.noreply.github.com>
Date: Thu, 27 Feb 2020 11:24:59 +1100
Subject: [PATCH] Add region argument (#7)

---
 README.md       |  2 ++
 cmd/main.go     |  5 +++--
 core/support.go | 48 +++++++++++++++++++++++++++++++++++++++++++++---
 core/types.go   |  1 +
 go.mod          |  2 ++
 5 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index dff119d..10ccc9f 100644
--- a/README.md
+++ b/README.md
@@ -57,6 +57,8 @@ Usage of ./bin/aws-limits-exporter:
         If non-empty, write log files in this directory
   -logtostderr
         log to standard error instead of files
+  -region string
+        The AWS region to show metrics for (default all regions).
   -stderrthreshold value
         logs at or above this threshold go to stderr
   -v value
diff --git a/cmd/main.go b/cmd/main.go
index 8137bdf..0048643 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -16,7 +16,8 @@ var (
 	// VERSION set by build script
 	VERSION = "UNKNOWN"
 
-	addr = flag.String("listen-address", ":8080", "The address to listen on for HTTP requests.")
+	addr   = flag.String("listen-address", ":8080", "The address to listen on for HTTP requests.")
+	region = flag.String("region", "", "The AWS region to show metrics for (default all regions).")
 )
 
 func main() {
@@ -24,7 +25,7 @@ func main() {
 
 	glog.Infof("AWS Limits Exporter v%s started.", VERSION)
 
-	exporter := core.NewSupportExporter()
+	exporter := core.NewSupportExporter(*region)
 	go exporter.RequestServiceLimitsRefreshLoop()
 
 	prometheus.Register(exporter)
diff --git a/core/support.go b/core/support.go
index af5d2c0..99cc48c 100644
--- a/core/support.go
+++ b/core/support.go
@@ -9,6 +9,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
+	"github.com/aws/aws-sdk-go/aws/endpoints"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/support"
 	"github.com/golang/glog"
@@ -117,6 +118,21 @@ var (
 	}
 )
 
+// validateRegionName ...
+func validateRegionName(region string) {
+	if region != "" {
+		availableRegions := endpoints.AwsPartition().Regions()
+		if _, ok := availableRegions[region]; !ok {
+			regions := make([]string, 0, len(availableRegions))
+			for key := range availableRegions {
+				regions = append(regions, key)
+			}
+
+			glog.Fatalf("Invalid AWS region %s, valid regions: %s", region, strings.Join(regions, ","))
+		}
+	}
+}
+
 // NewSupportClient ...
 func NewSupportClient() *SupportClientImpl {
 	creds := credentials.NewChainCredentials(
@@ -181,16 +197,29 @@ func (client *SupportClientImpl) DescribeServiceLimitsCheckResult(checkID string
 }
 
 // NewSupportExporter ...
-func NewSupportExporter() *SupportExporter {
+func NewSupportExporter(region string) *SupportExporter {
+	validateRegionName(region)
+
 	client := NewSupportClient()
 
 	return &SupportExporter{
 		supportClient: client,
+		metricsRegion: region,
 		metricsUsed:   map[string]*prometheus.Desc{},
 		metricsLimit:  map[string]*prometheus.Desc{},
 	}
 }
 
+// validateMetricRegion
+func (e *SupportExporter) validateMetricRegion(region string) bool {
+	if e.metricsRegion == "" {
+		return true
+	} else if e.metricsRegion == region {
+		return true
+	}
+	return false
+}
+
 // RequestServiceLimitsRefreshLoop ...
 func (e *SupportExporter) RequestServiceLimitsRefreshLoop() {
 	e.supportClient.RequestServiceLimitsRefreshLoop()
@@ -212,10 +241,18 @@ func (e *SupportExporter) Describe(ch chan<- *prometheus.Desc) {
 				continue
 			}
 
+			region := aws.StringValue(resource.Metadata[0])
+
+			// Filter region
+			if !e.validateMetricRegion(region) {
+				continue
+			}
+
 			serviceName := aws.StringValue(resource.Metadata[1])
 			serviceNameLower := strings.ToLower(serviceName)
-			e.metricsUsed[resourceID] = newServerMetric(aws.StringValue(resource.Metadata[0]), serviceNameLower, "used_total", "Current used amount of the given resource.", []string{"resource"})
-			e.metricsLimit[resourceID] = newServerMetric(aws.StringValue(resource.Metadata[0]), serviceNameLower, "limit_total", "Current limit of the given resource.", []string{"resource"})
+
+			e.metricsUsed[resourceID] = newServerMetric(region, serviceNameLower, "used_total", "Current used amount of the given resource.", []string{"resource"})
+			e.metricsLimit[resourceID] = newServerMetric(region, serviceNameLower, "limit_total", "Current limit of the given resource.", []string{"resource"})
 
 			ch <- e.metricsUsed[resourceID]
 			ch <- e.metricsLimit[resourceID]
@@ -240,6 +277,11 @@ func (e *SupportExporter) Collect(ch chan<- prometheus.Metric) {
 				continue
 			}
 
+			// Filter region
+			if !e.validateMetricRegion(aws.StringValue(resource.Metadata[0])) {
+				continue
+			}
+
 			resourceName := aws.StringValue(resource.Metadata[2])
 
 			metricLimit := e.metricsLimit[resourceID]
diff --git a/core/types.go b/core/types.go
index f275fef..55a28ca 100644
--- a/core/types.go
+++ b/core/types.go
@@ -20,6 +20,7 @@ type SupportClient interface {
 // SupportExporter ...
 type SupportExporter struct {
 	supportClient SupportClient
+	metricsRegion string
 	metricsUsed   map[string]*prometheus.Desc
 	metricsLimit  map[string]*prometheus.Desc
 }
diff --git a/go.mod b/go.mod
index 41f6d01..c7c7885 100644
--- a/go.mod
+++ b/go.mod
@@ -5,5 +5,7 @@ go 1.12
 require (
 	github.com/aws/aws-sdk-go v1.25.34
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.1 // indirect
 	github.com/prometheus/client_golang v1.2.1
 )