Add UnmarshalYAML for DropPolicy type

policy/policy.go

@@ -37,8 +37,9 @@ var (
 	// DefaultPolicy represents a default policy.
 	DefaultPolicy Policy
 
-	errNilPolicyProto      = errors.New("nil policy proto")
-	errInvalidPolicyString = errors.New("invalid policy string")
+	errNilPolicyProto          = errors.New("nil policy proto")
+	errInvalidPolicyString     = errors.New("invalid policy string")
+	errInvalidDropPolicyString = errors.New("invalid drop policy string")
 )
 
 // Policy contains a storage policy and a list of custom aggregation types.
@@ -202,3 +203,36 @@ func (p Policies) Equals(other Policies) bool {
 	}
 	return true
 }
+
+// UnmarshalYAML unmarshals a drop policy value from a string.
+func (p *DropPolicy) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var str string
+	if err := unmarshal(&str); err != nil {
+		return err
+	}
+
+	// Allow default string value (not specified) to mean default
+	if str == "" {
+		*p = DefaultDropPolicy
+		return nil
+	}
+
+	parsed, err := ParseDropPolicy(str)
+	if err != nil {
+		return err
+	}
+
+	*p = parsed
+	return nil
+}
+
+// ParseDropPolicy parses a drop policy.
+func ParseDropPolicy(str string) (DropPolicy, error) {
+	for _, valid := range validDropPolicies {
+		if valid.String() == str {
+			return valid, nil
+		}
+	}
+
+	return DefaultDropPolicy, errInvalidDropPolicyString
+}

policy/policy_test.go

@@ -229,3 +229,43 @@ func TestParsePolicyIntoProto(t *testing.T) {
 		require.Equal(t, input.expected, sp, input.str)
 	}
 }
+
+func TestDropPolicyUnmarshalYAML(t *testing.T) {
+	inputs := []struct {
+		str      string
+		expected DropPolicy
+	}{
+		{
+			str:      "",
+			expected: DropNone,
+		},
+		{
+			str:      DropNone.String(),
+			expected: DropNone,
+		},
+		{
+			str:      DropMust.String(),
+			expected: DropMust,
+		},
+		{
+			str:      DropIfOnlyMatch.String(),
+			expected: DropIfOnlyMatch,
+		},
+	}
+	for _, input := range inputs {
+		var p DropPolicy
+		require.NoError(t, yaml.Unmarshal([]byte(input.str), &p))
+		require.Equal(t, input.expected, p)
+	}
+}
+
+func TestDropPolicyUnmarshalYAMLErrors(t *testing.T) {
+	inputs := []string{
+		"drop_musty",
+		"drop_unknown",
+	}
+	for _, input := range inputs {
+		var p DropPolicy
+		require.Error(t, yaml.Unmarshal([]byte(input), &p))
+	}
+}