unmarshalling ta config with mapstructure decoder

Author: davidhaja

cmd/otel-allocator/config/config.go

@@ -21,21 +21,24 @@ import (
 	"fmt"
 	"io/fs"
 	"os"
-	"regexp"
+	"reflect"
 	"time"
 
 	"github.com/go-logr/logr"
 	"github.com/prometheus/common/model"
 	promconfig "github.com/prometheus/prometheus/config"
 	_ "github.com/prometheus/prometheus/discovery/install"
 	"github.com/spf13/pflag"
+
 	"gopkg.in/yaml.v2"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
 	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
+
+	"github.com/go-viper/mapstructure/v2"
 )
 
 const (
@@ -81,6 +84,101 @@ type HTTPSServerConfig struct {
 	TLSKeyFilePath  string `yaml:"tls_key_file_path,omitempty"`
 }
 
+// StringToModelDurationHookFunc returns a DecodeHookFuncType
+// that converts string to time.Duration, which can be used
+// as model.Duration.
+func StringToModelDurationHookFunc() mapstructure.DecodeHookFuncType {
+	return func(
+		f reflect.Type,
+		t reflect.Type,
+		data interface{},
+	) (interface{}, error) {
+		if f.Kind() != reflect.String {
+			return data, nil
+		}
+
+		if t != reflect.TypeOf(model.Duration(5)) {
+			return data, nil
+		}
+
+		return time.ParseDuration(data.(string))
+	}
+}
+
+// MapToPromConfig returns a DecodeHookFuncType that provides a mechanism
+// for decoding promconfig.Config involving its own unmarshal logic.
+func MapToPromConfig() mapstructure.DecodeHookFuncType {
+	return func(
+		f reflect.Type,
+		t reflect.Type,
+		data interface{},
+	) (interface{}, error) {
+		if f.Kind() != reflect.Map {
+			return data, nil
+		}
+
+		if t != reflect.TypeOf(&promconfig.Config{}) {
+			return data, nil
+		}
+
+		pConfig := &promconfig.Config{}
+
+		mb, err := yaml.Marshal(data.(map[any]any))
+		if err != nil {
+			return nil, err
+		}
+
+		err = yaml.Unmarshal(mb, pConfig)
+		if err != nil {
+			return nil, err
+		}
+		return pConfig, nil
+	}
+}
+
+// MapToLabelSelector returns a DecodeHookFuncType that
+// provides a mechanism for decoding both matchLabels and matchExpressions from camelcase to lowercase
+// because we use yaml unmarshaling that supports lowercase field names if no `yaml` tag is defined
+// and metav1.LabelSelector uses `json` tags.
+// If both the camelcase and lowercase version is present, then the camelcase version takes precedence.
+func MapToLabelSelector() mapstructure.DecodeHookFuncType {
+	return func(
+		f reflect.Type,
+		t reflect.Type,
+		data interface{},
+	) (interface{}, error) {
+		if f.Kind() != reflect.Map {
+			return data, nil
+		}
+
+		if t != reflect.TypeOf(&metav1.LabelSelector{}) {
+			return data, nil
+		}
+
+		result := &metav1.LabelSelector{}
+		fMap := data.(map[any]any)
+		if matchLabels, ok := fMap["matchLabels"]; ok {
+			fMap["matchlabels"] = matchLabels
+			delete(fMap, "matchLabels")
+		}
+		if matchExpressions, ok := fMap["matchExpressions"]; ok {
+			fMap["matchexpressions"] = matchExpressions
+			delete(fMap, "matchExpressions")
+		}
+
+		b, err := yaml.Marshal(fMap)
+		if err != nil {
+			return nil, err
+		}
+
+		err = yaml.Unmarshal(b, result)
+		if err != nil {
+			return nil, err
+		}
+		return result, nil
+	}
+}
+
 func LoadFromFile(file string, target *Config) error {
 	return unmarshal(target, file)
 }
@@ -154,23 +252,40 @@ func LoadFromCLI(target *Config, flagSet *pflag.FlagSet) error {
 	return nil
 }
 
+// unmarshal decodes the contents of the configFile into the cfg argument, using a
+// mapstructure decoder with the following notable behaviors.
+// Decodes time.Duration from strings (see StringToModelDurationHookFunc).
+// Allows custom unmarshaling for promconfig.Config struct that implements yaml.Unmarshaler (see MapToPromConfig).
+// Allows custom unmarshaling for metav1.LabelSelector struct using both camelcase and lowercase field names (see MapToLabelSelector).
 func unmarshal(cfg *Config, configFile string) error {
 	yamlFile, err := os.ReadFile(configFile)
 	if err != nil {
 		return err
 	}
 
-	// Changing matchLabels and matchExpressions from camel case to lower case
-	// because we use yaml unmarshaling that supports lower case field names if no `yaml` tag is defined
-	// and metav1.LabelSelector uses `json` tags.
-	reLabels := regexp.MustCompile(`([ \t\f\v]*)matchLabels:([ \t\f\v]*\n)`)
-	yamlFile = reLabels.ReplaceAll(yamlFile, []byte("${1}matchlabels:${2}"))
-	reExpressions := regexp.MustCompile(`([ \t\f\v]*)matchExpressions:([ \t\f\v]*\n)`)
-	yamlFile = reExpressions.ReplaceAll(yamlFile, []byte("${1}matchexpressions:${2}"))
-
-	if err = yaml.Unmarshal(yamlFile, cfg); err != nil {
+	m := make(map[string]interface{})
+	if err := yaml.Unmarshal(yamlFile, &m); err != nil {
 		return fmt.Errorf("error unmarshaling YAML: %w", err)
 	}
+
+	dc := mapstructure.DecoderConfig{
+		TagName: "yaml",
+		Result:  cfg,
+		DecodeHook: mapstructure.ComposeDecodeHookFunc(
+			StringToModelDurationHookFunc(),
+			MapToPromConfig(),
+			MapToLabelSelector(),
+		),
+	}
+
+	decoder, err := mapstructure.NewDecoder(&dc)
+	if err != nil {
+		return err
+	}
+	if err := decoder.Decode(m); err != nil {
+		return err
+	}
+
 	return nil
 }
 

go.mod

@@ -228,3 +228,5 @@ require (
 	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
 )
+
+require github.com/go-viper/mapstructure/v2 v2.2.1

go.sum

@@ -236,6 +236,8 @@ github.com/go-resty/resty/v2 v2.13.1/go.mod h1:GznXlLxkq6Nh4sU59rPmUw3VtgpO3aS96
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI=
 github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8=
+github.com/go-viper/mapstructure/v2 v2.2.1 h1:ZAaOCxANMuZx5RCeg0mBdEZk7DZasvvZIxtHqx8aGss=
+github.com/go-viper/mapstructure/v2 v2.2.1/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
 github.com/go-zookeeper/zk v1.0.3 h1:7M2kwOsc//9VeeFiPtf+uSJlVpU66x9Ba5+8XK7/TDg=
 github.com/go-zookeeper/zk v1.0.3/go.mod h1:nOB03cncLtlp4t+UAkGSV+9beXP/akpekBwL+UX1Qcw=
 github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=