model.ValidationScheme: Support encoding as YAML (#799)

* model.ValidationScheme: Support encoding as YAML

---------

Signed-off-by: Arve Knudsen <[email protected]>

Author: aknuds1

model/metric.go

@@ -24,6 +24,7 @@ import (
 
 	dto "github.com/prometheus/client_model/go"
 	"google.golang.org/protobuf/proto"
+	"gopkg.in/yaml.v2"
 )
 
 var (
@@ -62,16 +63,70 @@ var (
 type ValidationScheme int
 
 const (
+	// UnsetValidation represents an undefined ValidationScheme.
+	// Should not be used in practice.
+	UnsetValidation ValidationScheme = iota
+
 	// LegacyValidation is a setting that requires that all metric and label names
 	// conform to the original Prometheus character requirements described by
 	// MetricNameRE and LabelNameRE.
-	LegacyValidation ValidationScheme = iota
+	LegacyValidation
 
 	// UTF8Validation only requires that metric and label names be valid UTF-8
 	// strings.
 	UTF8Validation
 )
 
+var (
+	_ yaml.Marshaler = UnsetValidation
+	_ fmt.Stringer   = UnsetValidation
+)
+
+// String returns the string representation of s.
+func (s ValidationScheme) String() string {
+	switch s {
+	case UnsetValidation:
+		return "unset"
+	case LegacyValidation:
+		return "legacy"
+	case UTF8Validation:
+		return "utf8"
+	default:
+		panic(fmt.Errorf("unhandled ValidationScheme: %d", s))
+	}
+}
+
+// MarshalYAML implements the yaml.Marshaler interface.
+func (s ValidationScheme) MarshalYAML() (any, error) {
+	switch s {
+	case UnsetValidation:
+		return "", nil
+	case LegacyValidation, UTF8Validation:
+		return s.String(), nil
+	default:
+		panic(fmt.Errorf("unhandled ValidationScheme: %d", s))
+	}
+}
+
+// UnmarshalYAML implements the yaml.Unmarshaler interface.
+func (s *ValidationScheme) UnmarshalYAML(unmarshal func(any) error) error {
+	var scheme string
+	if err := unmarshal(&scheme); err != nil {
+		return err
+	}
+	switch scheme {
+	case "":
+		// Don't change the value.
+	case "legacy":
+		*s = LegacyValidation
+	case "utf8":
+		*s = UTF8Validation
+	default:
+		return fmt.Errorf("unrecognized ValidationScheme: %q", scheme)
+	}
+	return nil
+}
+
 type EscapingScheme int
 
 const (
@@ -185,7 +240,7 @@ func IsValidMetricName(n LabelValue) bool {
 		}
 		return utf8.ValidString(string(n))
 	default:
-		panic(fmt.Sprintf("Invalid name validation scheme requested: %d", NameValidationScheme))
+		panic(fmt.Sprintf("Invalid name validation scheme requested: %s", NameValidationScheme.String()))
 	}
 }
 

model/metric_test.go

@@ -14,12 +14,16 @@
 package model
 
 import (
+	"errors"
+	"strings"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	dto "github.com/prometheus/client_model/go"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
+	"gopkg.in/yaml.v2"
 )
 
 func testMetric(t testing.TB) {
@@ -89,6 +93,115 @@ func BenchmarkMetric(b *testing.B) {
 	}
 }
 
+func TestValidationScheme(t *testing.T) {
+	var scheme ValidationScheme
+	require.Equal(t, UnsetValidation, scheme)
+}
+
+func TestValidationScheme_String(t *testing.T) {
+	for _, tc := range []struct {
+		name   string
+		scheme ValidationScheme
+		want   string
+	}{
+		{
+			name:   "Unset",
+			scheme: UnsetValidation,
+			want:   "unset",
+		},
+		{
+			name:   "Legacy",
+			scheme: LegacyValidation,
+			want:   "legacy",
+		},
+		{
+			name:   "UTF8",
+			scheme: UTF8Validation,
+			want:   "utf8",
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.want, tc.scheme.String())
+		})
+	}
+}
+
+func TestValidationScheme_MarshalYAML(t *testing.T) {
+	for _, tc := range []struct {
+		name   string
+		scheme ValidationScheme
+		want   string
+	}{
+		{
+			name:   "Unset",
+			scheme: UnsetValidation,
+			want:   `""`,
+		},
+		{
+			name:   "Legacy",
+			scheme: LegacyValidation,
+			want:   "legacy",
+		},
+		{
+			name:   "UTF8",
+			scheme: UTF8Validation,
+			want:   "utf8",
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			marshaled, err := yaml.Marshal(tc.scheme)
+			require.NoError(t, err)
+			require.Equal(t, tc.want, strings.TrimSpace(string(marshaled)))
+		})
+	}
+}
+
+func TestValidationScheme_UnmarshalYAML(t *testing.T) {
+	for _, tc := range []struct {
+		name      string
+		input     string
+		want      ValidationScheme
+		wantError error
+	}{
+		{
+			name:  "Unset empty input",
+			input: "",
+			want:  UnsetValidation,
+		},
+		{
+			name:  "Unset quoted input",
+			input: `""`,
+			want:  UnsetValidation,
+		},
+		{
+			name:  "Legacy",
+			input: "legacy",
+			want:  LegacyValidation,
+		},
+		{
+			name:  "UTF8",
+			input: "utf8",
+			want:  UTF8Validation,
+		},
+		{
+			name:      "Invalid",
+			input:     "invalid",
+			wantError: errors.New(`unrecognized ValidationScheme: "invalid"`),
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			scheme := UnsetValidation
+			err := yaml.Unmarshal([]byte(tc.input), &scheme)
+			if tc.wantError == nil {
+				require.NoError(t, err)
+				require.Equal(t, tc.want, scheme)
+			} else {
+				require.EqualError(t, err, tc.wantError.Error())
+			}
+		})
+	}
+}
+
 func TestMetricNameIsLegacyValid(t *testing.T) {
 	scenarios := []struct {
 		mn          LabelValue