Middleware: Add team metadata to HTTP handlers (#71010)

Signed-off-by: bergquist <carl.bergquist@gmail.com>

pkg/middleware/request_metadata_test.go

@@ -0,0 +1,40 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/middleware/requestmeta"
+	"github.com/grafana/grafana/pkg/web"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRequestMetaDefault(t *testing.T) {
+	m := web.New()
+	m.Use(requestmeta.SetupRequestMetadata())
+
+	m.Get("/", func(rw http.ResponseWriter, req *http.Request) {
+		v := requestmeta.GetRequestMetaData(req.Context())
+		assert.Equal(t, requestmeta.TeamCore, v.Team)
+	})
+
+	req, _ := http.NewRequest(http.MethodGet, "/", nil)
+	m.ServeHTTP(httptest.NewRecorder(), req)
+}
+
+func TestRequestMetaNewTeam(t *testing.T) {
+	m := web.New()
+	m.Use(requestmeta.SetupRequestMetadata())
+
+	m.Get("/",
+		requestmeta.SetOwner(requestmeta.TeamAlerting), // set new owner for this route.
+		func(rw http.ResponseWriter, req *http.Request) {
+			v := requestmeta.GetRequestMetaData(req.Context())
+			assert.Equal(t, requestmeta.TeamAlerting, v.Team)
+		})
+
+	r, err := http.NewRequest(http.MethodGet, "/", nil)
+	assert.NoError(t, err)
+	m.ServeHTTP(httptest.NewRecorder(), r)
+}

pkg/middleware/request_metrics.go

@@ -11,21 +11,23 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/metrics"
 	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/middleware/requestmeta"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/web"
 )
 
 var (
-	httpRequestsInFlight         prometheus.Gauge
-	httpRequestDurationHistogram *prometheus.HistogramVec
-
 	// DefBuckets are histogram buckets for the response time (in seconds)
 	// of a network service, including one that is responding very slowly.
 	defBuckets = []float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10, 25}
 )
 
-func init() {
-	httpRequestsInFlight = prometheus.NewGauge(
+// RequestMetrics is a middleware handler that instruments the request.
+func RequestMetrics(features featuremgmt.FeatureToggles, cfg *setting.Cfg, promRegister prometheus.Registerer) web.Middleware {
+	log := log.New("middleware.request-metrics")
+
+	httpRequestsInFlight := prometheus.NewGauge(
 		prometheus.GaugeOpts{
 			Namespace: "grafana",
 			Name:      "http_request_in_flight",
@@ -33,22 +35,22 @@ func init() {
 		},
 	)
 
-	httpRequestDurationHistogram = prometheus.NewHistogramVec(
+	histogramLabels := []string{"handler", "status_code", "method"}
+	if cfg.MetricsIncludeTeamLabel {
+		histogramLabels = append(histogramLabels, "team")
+	}
+
+	httpRequestDurationHistogram := prometheus.NewHistogramVec(
 		prometheus.HistogramOpts{
 			Namespace: "grafana",
 			Name:      "http_request_duration_seconds",
 			Help:      "Histogram of latencies for HTTP requests.",
 			Buckets:   defBuckets,
 		},
-		[]string{"handler", "status_code", "method"},
+		histogramLabels,
 	)
 
-	prometheus.MustRegister(httpRequestsInFlight, httpRequestDurationHistogram)
-}
-
-// RequestMetrics is a middleware handler that instruments the request.
-func RequestMetrics(features featuremgmt.FeatureToggles) web.Middleware {
-	log := log.New("middleware.request-metrics")
+	promRegister.MustRegister(httpRequestsInFlight, httpRequestDurationHistogram)
 
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -77,10 +79,17 @@ func RequestMetrics(features featuremgmt.FeatureToggles) web.Middleware {
 				}
 			}
 
+			labelValues := []string{handler, code, r.Method}
+			if cfg.MetricsIncludeTeamLabel {
+				rmd := requestmeta.GetRequestMetaData(r.Context())
+				labelValues = append(labelValues, rmd.Team)
+			}
+
 			// avoiding the sanitize functions for in the new instrumentation
 			// since they dont make much sense. We should remove them later.
 			histogram := httpRequestDurationHistogram.
-				WithLabelValues(handler, code, r.Method)
+				WithLabelValues(labelValues...)
+
 			if traceID := tracing.TraceIDFromContext(r.Context(), true); traceID != "" {
 				// Need to type-convert the Observer to an
 				// ExemplarObserver. This will always work for a

pkg/middleware/requestmeta/request_metadata.go

@@ -0,0 +1,75 @@
+package requestmeta
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/grafana/grafana/pkg/web"
+)
+
+const (
+	TeamAlerting = "alerting"
+	TeamAuth     = "auth"
+	TeamCore     = "core"
+)
+
+type rMDContextKey struct{}
+
+type RequestMetaData struct {
+	Team string
+}
+
+var requestMetaDataContextKey = rMDContextKey{}
+
+// SetupRequestMetadata injects defaul request metadata values
+// on the request context.
+func SetupRequestMetadata() web.Middleware {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			rmd := defaultRequestMetadata()
+
+			ctx := context.WithValue(r.Context(), requestMetaDataContextKey, rmd)
+			*r = *r.WithContext(ctx)
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
+// GetRequestMetaData returns the request metadata for the context.
+// if request metadata is missing it will return the default values.
+func GetRequestMetaData(ctx context.Context) *RequestMetaData {
+	val := ctx.Value(requestMetaDataContextKey)
+
+	value, ok := val.(*RequestMetaData)
+	if ok {
+		return value
+	}
+
+	return defaultRequestMetadata()
+}
+
+// SetRequestMetaData returns an `web.Handler` that overrides the request metadata
+// with the provided param.
+func SetRequestMetaData(rmd RequestMetaData) web.Handler {
+	return func(w http.ResponseWriter, r *http.Request) {
+		v := GetRequestMetaData(r.Context())
+		if rmd.Team != "" {
+			v.Team = rmd.Team
+		}
+	}
+}
+
+// SetOwner returns an `web.Handler` that sets the team name for an request.
+func SetOwner(team string) web.Handler {
+	return func(w http.ResponseWriter, r *http.Request) {
+		v := GetRequestMetaData(r.Context())
+		v.Team = team
+	}
+}
+
+func defaultRequestMetadata() *RequestMetaData {
+	return &RequestMetaData{
+		Team: TeamCore,
+	}
+}