Add support for mixin alerting (#16)

Author: nlamirault

generate/csv.go

@@ -4,8 +4,8 @@ import "fmt"
 
 // CSV finds all rules at the given path and its
 // subdirectories and generates a CSV file of the found rules.
-func CSV(path string) (string, error) {
-	ruleGroups, err := getRuleGroups(path)
+func CSV(path string, input string) (string, error) {
+	ruleGroups, err := getRuleGroups(path, input)
 	if err != nil {
 		return "", fmt.Errorf("get rule groups: %w", err)
 	}

generate/csv_test.go

@@ -11,7 +11,7 @@ func TestCSV(t *testing.T) {
 		t.Fatal("read:", err)
 	}
 
-	actual, err := Generate("../examples", ".csv")
+	actual, err := Generate("../examples", ".csv", "kubernetes")
 	if err != nil {
 		t.Fatal("generate:", err)
 	}

generate/generate.go

@@ -3,100 +3,44 @@ package generate
 import (
 	"errors"
 	"fmt"
-	"io/ioutil"
 	"os"
 	"path/filepath"
 	"regexp"
 	"sort"
-
-	"github.com/ghodss/yaml"
 )
 
-type typeMeta struct {
-	Kind string `json:"kind"`
-}
-
-type prometheusRule struct {
-	Spec prometheusRuleSpec `json:"spec"`
-}
-
-type prometheusRuleSpec struct {
-	Groups []ruleGroup `json:"groups"`
-}
-
-type ruleGroup struct {
-	Name  string `json:"name"`
-	Rules []rule `json:"rules"`
-}
-
-type rule struct {
-	Alert       string            `json:"alert"`
-	Labels      map[string]string `json:"labels"`
-	Annotations map[string]string `json:"annotations"`
-}
-
 // Generate finds all rules at the given path and its
 // subdirectories and generates documentation with the
 // specified file extension, including the period.
-func Generate(path string, output string) (string, error) {
+func Generate(path string, output string, input string) (string, error) {
 	switch output {
 	case ".md":
-		return Markdown(path)
+		return Markdown(path, input)
 	case ".csv":
-		return CSV(path)
+		return CSV(path, input)
 	default:
 		return "", errors.New("output format not supported")
 	}
 }
 
-func getRuleGroups(path string) ([]ruleGroup, error) {
+func getRuleGroups(path string, input string) ([]ruleGroup, error) {
 	files, err := getYamlFiles(path)
 	if err != nil {
 		return nil, fmt.Errorf("get yaml files: %w", err)
 	}
 
 	var alertGroups []ruleGroup
-	for _, file := range files {
-		fileBytes, err := ioutil.ReadFile(file)
+	if "kubernetes" == input {
+		alertGroups, err = getKubernetesRuleGroups(files)
 		if err != nil {
-			return nil, fmt.Errorf("open file: %w", err)
+			return nil, err
 		}
-
-		var typeMeta typeMeta
-		if err := yaml.Unmarshal(fileBytes, &typeMeta); err != nil {
-			continue
-		}
-		if typeMeta.Kind != "PrometheusRule" {
-			continue
-		}
-
-		var prometheusRule prometheusRule
-		if err := yaml.Unmarshal(fileBytes, &prometheusRule); err != nil {
-			continue
-		}
-
-		for _, group := range prometheusRule.Spec.Groups {
-			var alertRules []rule
-			for _, rule := range group.Rules {
-				if rule.Alert == "" {
-					continue
-				}
-
-				alertRules = append(alertRules, rule)
-			}
-
-			if len(alertRules) == 0 {
-				continue
-			}
-
-			alertGroup := ruleGroup{
-				Name:  group.Name,
-				Rules: alertRules,
-			}
-			alertGroups = append(alertGroups, alertGroup)
+	} else {
+		alertGroups, err = getMixinRuleGroups(files)
+		if err != nil {
+			return nil, err
 		}
 	}
-
 	sort.Slice(alertGroups, func(i int, j int) bool {
 		return alertGroups[i].Name < alertGroups[j].Name
 	})

generate/markdown.go

@@ -8,8 +8,8 @@ import (
 // Markdown finds all rules at the given path and its
 // subdirectories and generates Markdown documentation
 // from the found rules.
-func Markdown(path string) (string, error) {
-	ruleGroups, err := getRuleGroups(path)
+func Markdown(path string, input string) (string, error) {
+	ruleGroups, err := getRuleGroups(path, input)
 	if err != nil {
 		return "", fmt.Errorf("get rule groups: %w", err)
 	}

generate/markdown_test.go

@@ -11,7 +11,7 @@ func TestMarkdown(t *testing.T) {
 		t.Fatal("read:", err)
 	}
 
-	actual, err := Generate("../examples", ".md")
+	actual, err := Generate("../examples", ".md", "kubernetes")
 	if err != nil {
 		t.Fatal("generate:", err)
 	}

generate/prometheus.go

@@ -0,0 +1,114 @@
+package generate
+
+import (
+	"fmt"
+	"io/ioutil"
+
+	"github.com/ghodss/yaml"
+)
+
+type typeMeta struct {
+	Kind string `json:"kind"`
+}
+
+type prometheusRule struct {
+	Spec prometheusRuleSpec `json:"spec"`
+}
+
+type prometheusRuleSpec struct {
+	Groups []ruleGroup `json:"groups"`
+}
+
+type ruleGroup struct {
+	Name  string `json:"name"`
+	Rules []rule `json:"rules"`
+}
+
+type rule struct {
+	Alert       string            `json:"alert"`
+	Labels      map[string]string `json:"labels"`
+	Annotations map[string]string `json:"annotations"`
+}
+
+func getMixinRuleGroups(files []string) ([]ruleGroup, error) {
+	var alertGroups []ruleGroup
+
+	for _, file := range files {
+		fileBytes, err := ioutil.ReadFile(file)
+		if err != nil {
+			return nil, fmt.Errorf("open file: %w", err)
+		}
+
+		var groups prometheusRuleSpec
+		if err := yaml.Unmarshal(fileBytes, &groups); err != nil {
+			continue
+		}
+
+		for _, group := range groups.Groups {
+			alertGroup, err := extractGroupAlerts(group)
+			if err != nil {
+				return nil, err
+			}
+			alertGroups = append(alertGroups, *alertGroup)
+		}
+
+	}
+
+	return alertGroups, nil
+}
+
+func getKubernetesRuleGroups(files []string) ([]ruleGroup, error) {
+	var alertGroups []ruleGroup
+	for _, file := range files {
+		fileBytes, err := ioutil.ReadFile(file)
+		if err != nil {
+			return nil, fmt.Errorf("open file: %w", err)
+		}
+
+		var typeMeta typeMeta
+		if err := yaml.Unmarshal(fileBytes, &typeMeta); err != nil {
+			continue
+		}
+		if typeMeta.Kind != "PrometheusRule" {
+			continue
+		}
+
+		var prometheusRule prometheusRule
+		if err := yaml.Unmarshal(fileBytes, &prometheusRule); err != nil {
+			continue
+		}
+
+		for _, group := range prometheusRule.Spec.Groups {
+			alertGroup, err := extractGroupAlerts(group)
+			if err != nil {
+				return nil, err
+			}
+			if alertGroup != nil {
+				alertGroups = append(alertGroups, *alertGroup)
+			}
+		}
+
+	}
+	return alertGroups, nil
+}
+
+func extractGroupAlerts(group ruleGroup) (*ruleGroup, error) {
+	var alertRules []rule
+	for _, rule := range group.Rules {
+		if rule.Alert == "" {
+			return nil, nil
+		}
+
+		alertRules = append(alertRules, rule)
+	}
+
+	if len(alertRules) == 0 {
+		return nil, nil
+	}
+
+	alertGroup := ruleGroup{
+		Name:  group.Name,
+		Rules: alertRules,
+	}
+	return &alertGroup, nil
+}

internal/commands/generate.go

@@ -22,6 +22,10 @@ func NewGenerateCommand() *cobra.Command {
 			if err := viper.BindPFlag("out", cmd.Flags().Lookup("out")); err != nil {
 				return fmt.Errorf("bind out flag: %w", err)
 			}
+			if err := viper.BindPFlag("in", cmd.Flags().Lookup("in")); err != nil {
+				return fmt.Errorf("bind in flag: %w", err)
+			}
+
 			return nil
 		},
 
@@ -40,7 +44,7 @@ func NewGenerateCommand() *cobra.Command {
 	}
 
 	cmd.Flags().StringP("out", "o", "alerts.md", "File name or directory for the alert documentation")
-
+	cmd.Flags().StringP("in", "i", "kubernetes", "Alert style: Kubernetes or Mixin")
 	return &cmd
 }
 
@@ -50,7 +54,9 @@ func runGenerateCommand(path string) error {
 		outputPath = filepath.Join(outputPath, "alerts.md")
 	}
 
-	output, err := generate.Generate(path, filepath.Ext(outputPath))
+	input := viper.GetString("in")
+
+	output, err := generate.Generate(path, filepath.Ext(outputPath), input)
 	if err != nil {
 		return fmt.Errorf("render: %w", err)
 	}