Auto-discover regions for /probe/metrics endpoint (#68)

* Auto-discover regions for /probe/metrics endpoint * discoverResourceRegions
webdevops · Dec 12, 2023 · fadc573 · fadc573
1 parent e0da459
commit fadc573
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -304,7 +304,7 @@ one metric request per subscription and region
 | GET parameter        | Default                   | Required | Multiple | Description                                                                                                                                          |
 |----------------------|---------------------------|----------|----------|------------------------------------------------------------------------------------------------------------------------------------------------------|
 | `subscription`       |                           | **yes**  | **yes**  | Azure Subscription ID                                                                                                                                |
-| `region`             |                           | **yes**  | **yes**  | Azure Regions (eg. `westeurope`, `northeurope`)                                                                                                      |
+| `region`             |                           | no       | **yes**  | Azure Regions (eg. `westeurope`, `northeurope`). If omit, ResourceGrapth will be used to discover regions                                            |
 | `resourceType`       |                           | **yes**  | no       | Azure Resource type                                                                                                                                  |
 | `timespan`           | `PT1M`                    | no       | no       | Metric timespan                                                                                                                                      |
 | `interval`           |                           | no       | no       | Metric timespan                                                                                                                                      |

diff --git a/metrics/prober.go b/metrics/prober.go
@@ -2,6 +2,7 @@ package metrics
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"strings"
 	"time"
@@ -181,9 +182,16 @@ func (p *MetricProber) collectMetricsFromSubscriptions() {
 	subscriptionIterator.SetConcurrency(p.Conf.Prober.ConcurrencySubscription)
 
 	go func() {
+		regions, err := p.discoverResourceRegions()
+		if err != nil {
+			p.logger.Error(fmt.Errorf("error getting subscription locations: %w", err))
+			return
+		}
+
+		err = subscriptionIterator.ForEachAsync(p.logger, func(subscription *armsubscriptions.Subscription, logger *zap.SugaredLogger) {
+			subscriptionRegions := regions[*subscription.SubscriptionID]
 
-		err := subscriptionIterator.ForEachAsync(p.logger, func(subscription *armsubscriptions.Subscription, logger *zap.SugaredLogger) {
-			for _, region := range p.settings.Regions {
+			for _, region := range subscriptionRegions {
 				client, err := p.MetricsClient(*subscription.SubscriptionID)
 				if err != nil {
 					// FIXME: find a better way to report errors
@@ -267,6 +275,38 @@ func (p *MetricProber) collectMetricsFromSubscriptions() {
 	}
 }
 
+func (p *MetricProber) discoverResourceRegions() (map[string][]string, error) {
+	regions := map[string][]string{}
+
+	for _, subscriptionId := range p.settings.Subscriptions {
+		if len(p.settings.Regions) == 0 {
+			regions[subscriptionId] = []string{}
+		} else {
+			regions[subscriptionId] = p.settings.Regions
+		}
+	}
+
+	if len(p.settings.Regions) != 0 {
+		return regions, nil
+	}
+
+	query := fmt.Sprintf(`Resources | where type == "%s" | summarize count() by subscriptionId, location`, strings.ToLower(p.settings.ResourceType))
+
+	results, err := p.AzureClient.ExecuteResourceGraphQuery(p.ctx, p.settings.Subscriptions, query)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, row := range results {
+		subscriptionId := row["subscriptionId"].(string)
+		location := row["location"].(string)
+
+		regions[subscriptionId] = append(regions[subscriptionId], location)
+	}
+
+	return regions, nil
+}
+
 func (p *MetricProber) collectMetricsFromTargets() {
 	metricsChannel := make(chan PrometheusMetricResult)
 

diff --git a/probe_metrics_subscription.go b/probe_metrics_subscription.go
@@ -47,12 +47,6 @@ func probeMetricsSubscriptionHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	if _, err = paramsGetListRequired(r.URL.Query(), "region"); err != nil {
-		contextLogger.Warnln(err)
-		http.Error(w, err.Error(), http.StatusBadRequest)
-		return
-	}
-
 	prober := metrics.NewMetricProber(ctx, contextLogger, w, &settings, opts)
 	prober.SetUserAgent(UserAgent + gitTag)
 	prober.SetAzureClient(AzureClient)