feature: add ResponseFormatIncludingMetrics helper func

As there is existing similar function for encoder `NegotiateIncludingOpenMetrics`, introducing function for parsing format from `content-type` header will be useful. Signed-off-by: Bart Smykla <[email protected]>
prometheus · Jul 13, 2023 · 06a8274 · 06a8274
1 parent 94bf982
commit 06a8274
Show file tree

Hide file tree

Showing 2 changed files with 117 additions and 0 deletions.
diff --git a/expfmt/decode.go b/expfmt/decode.go
@@ -23,6 +23,7 @@ import (
 	dto "github.com/prometheus/client_model/go"
 
 	"github.com/matttproud/golang_protobuf_extensions/pbutil"
+
 	"github.com/prometheus/common/model"
 )
 
@@ -69,6 +70,48 @@ func ResponseFormat(h http.Header) Format {
 	return FmtUnknown
 }
 
+// ResponseFormatIncludingOpenMetrics works like ResponseFormat but includes
+// FmtOpenMetrics as an option for the result. Note that this function is
+// temporary and will disappear once FmtOpenMetrics is fully supported and as
+// such may be returned by the normal ResponseFormat function.
+func ResponseFormatIncludingOpenMetrics(h http.Header) Format {
+	ct := h.Get(hdrContentType)
+
+	mediatype, params, err := mime.ParseMediaType(ct)
+	if err != nil {
+		return FmtUnknown
+	}
+
+	const textType = "text/plain"
+
+	switch mediatype {
+	case ProtoType:
+		if p, ok := params["proto"]; ok && p != ProtoProtocol {
+			return FmtUnknown
+		}
+		if e, ok := params["encoding"]; ok && e != "delimited" {
+			return FmtUnknown
+		}
+		return FmtProtoDelim
+
+	case textType:
+		if v, ok := params["version"]; ok && v != TextVersion {
+			return FmtUnknown
+		}
+		return FmtText
+
+	case OpenMetricsType:
+		switch params["version"] {
+		case OpenMetricsVersion_1_0_0:
+			return FmtOpenMetrics_1_0_0
+		case OpenMetricsVersion_0_0_1, "":
+			return FmtOpenMetrics_0_0_1
+		}
+	}
+
+	return FmtUnknown
+}
+
 // NewDecoder returns a new decoder based on the given input format.
 // If the input format does not imply otherwise, a text format decoder is returned.
 func NewDecoder(r io.Reader, format Format) Decoder {

diff --git a/expfmt/decode_test.go b/expfmt/decode_test.go
@@ -418,16 +418,90 @@ func testDiscriminatorHTTPHeader(t testing.TB) {
 	}
 }
 
+func testDiscriminatorHTTPHeaderOpenMetrics(t testing.TB) {
+	var scenarios = []struct {
+		input  map[string]string
+		output Format
+	}{
+		// OpenMetrics
+		{
+			input:  map[string]string{"Content-Type": `application/openmetrics-text`},
+			output: FmtOpenMetrics_0_0_1,
+		},
+		{
+			input:  map[string]string{"Content-Type": `application/openmetrics-text;version=0.0.1`},
+			output: FmtOpenMetrics_0_0_1,
+		},
+		{
+			input:  map[string]string{"Content-Type": `application/openmetrics-text;version=1.0.0`},
+			output: FmtOpenMetrics_1_0_0,
+		},
+		// Other
+		{
+			input:  map[string]string{"Content-Type": `application/vnd.google.protobuf; proto="io.prometheus.client.MetricFamily"; encoding="delimited"`},
+			output: FmtProtoDelim,
+		},
+		{
+			input:  map[string]string{"Content-Type": `application/vnd.google.protobuf; proto="illegal"; encoding="delimited"`},
+			output: FmtUnknown,
+		},
+		{
+			input:  map[string]string{"Content-Type": `application/vnd.google.protobuf; proto="io.prometheus.client.MetricFamily"; encoding="illegal"`},
+			output: FmtUnknown,
+		},
+		{
+			input:  map[string]string{"Content-Type": `text/plain; version=0.0.4`},
+			output: FmtText,
+		},
+		{
+			input:  map[string]string{"Content-Type": `text/plain`},
+			output: FmtText,
+		},
+		{
+			input:  map[string]string{"Content-Type": `text/plain; version=0.0.3`},
+			output: FmtUnknown,
+		},
+	}
+
+	for i, scenario := range scenarios {
+		var header http.Header
+
+		if len(scenario.input) > 0 {
+			header = http.Header{}
+		}
+
+		for key, value := range scenario.input {
+			header.Add(key, value)
+		}
+
+		actual := ResponseFormatIncludingOpenMetrics(header)
+
+		if scenario.output != actual {
+			t.Errorf("%d. expected %s, got %s", i, scenario.output, actual)
+		}
+	}
+}
+
 func TestDiscriminatorHTTPHeader(t *testing.T) {
 	testDiscriminatorHTTPHeader(t)
 }
 
+func TestDiscriminatorHTTPHeaderOpenMetrics(t *testing.T) {
+	testDiscriminatorHTTPHeaderOpenMetrics(t)
+}
+
 func BenchmarkDiscriminatorHTTPHeader(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		testDiscriminatorHTTPHeader(b)
 	}
 }
 
+func BenchmarkDiscriminatorHTTPHeaderOpenMetrics(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		testDiscriminatorHTTPHeaderOpenMetrics(b)
+	}
+}
+
 func TestExtractSamples(t *testing.T) {
 	var (
 		goodMetricFamily1 = &dto.MetricFamily{