Add id-labels flag (#3)

Author: Guðmundur Björn Birkisson

cmd/sloth/commands/generate.go

@@ -36,21 +36,23 @@ type generateCommand struct {
 	disableAlerts         bool
 	disableOptimizedRules bool
 	extraLabels           map[string]string
+	idLabels              map[string]string
 	sliPluginsPaths       []string
 	sloPeriodWindowsPath  string
 	sloPeriod             string
 }
 
 // NewGenerateCommand returns the generate command.
 func NewGenerateCommand(app *kingpin.Application) Command {
-	c := &generateCommand{extraLabels: map[string]string{}}
+	c := &generateCommand{extraLabels: map[string]string{}, idLabels: map[string]string{}}
 	cmd := app.Command("generate", "Generates Prometheus SLOs.")
 	cmd.Flag("input", "SLO spec input file path or directory (if directory is used, slos will be discovered recursively and out must be a directory).").Short('i').StringVar(&c.slosInput)
 	cmd.Flag("out", "Generated rules output file path or directory. If `-` it will use stdout (if input is a directory this must be a directory).").Default("-").Short('o').StringVar(&c.slosOut)
 	cmd.Flag("fs-exclude", "Filter regex to ignore matched discovered SLO file paths (used with directory based input/output).").Short('e').StringVar(&c.slosExcludeRegex)
 	cmd.Flag("fs-include", "Filter regex to include matched discovered SLO file paths, everything else will be ignored. Exclude has preference (used with directory based input/output).").Short('n').StringVar(&c.slosIncludeRegex)
 
 	cmd.Flag("extra-labels", "Extra labels that will be added to all the generated Prometheus rules ('key=value' form, can be repeated).").Short('l').StringMapVar(&c.extraLabels)
+	cmd.Flag("id-labels", "Id labels that used as filters for generated recording rules. These will also be added as extra labels ('key=value' form, can be repeated).").Short('d').StringMapVar(&c.idLabels)
 	cmd.Flag("disable-recordings", "Disables recording rules generation.").BoolVar(&c.disableRecordings)
 	cmd.Flag("disable-alerts", "Disables alert rules generation.").BoolVar(&c.disableAlerts)
 	cmd.Flag("sli-plugins-path", "The path to SLI plugins (can be repeated), if not set it disable plugins support.").Short('p').StringsVar(&c.sliPluginsPaths)
@@ -94,6 +96,11 @@ func (g generateCommand) Run(ctx context.Context, config RootConfig) error {
 		}
 	}
 
+	// Make sure id labels are set in extra labels as well
+	for key, value := range g.idLabels {
+		g.extraLabels[key] = value
+	}
+
 	// SLO period.
 	sp, err := prometheusmodel.ParseDuration(g.sloPeriod)
 	if err != nil {
@@ -245,6 +252,7 @@ func (g generateCommand) Run(ctx context.Context, config RootConfig) error {
 		disableAlerts:         g.disableAlerts,
 		disableOptimizedRules: g.disableOptimizedRules,
 		extraLabels:           g.extraLabels,
+		idLabels:              g.idLabels,
 	}
 
 	for _, genTarget := range genTargets {
@@ -305,6 +313,7 @@ type generator struct {
 	disableAlerts         bool
 	disableOptimizedRules bool
 	extraLabels           map[string]string
+	idLabels              map[string]string
 }
 
 // GeneratePrometheus generates the SLOs based on a raw regular Prometheus spec format input and outs a Prometheus raw yaml.
@@ -434,6 +443,7 @@ func (g generator) generateRules(ctx context.Context, info info.Info, slos prome
 
 	result, err := controller.Generate(ctx, generate.Request{
 		ExtraLabels: g.extraLabels,
+		IDLabels:    g.idLabels,
 		Info:        info,
 		SLOGroup:    slos,
 	})

cmd/sloth/commands/k8scontroller.go

@@ -53,6 +53,7 @@ const (
 
 type kubeControllerCommand struct {
 	extraLabels           map[string]string
+	idLabels              map[string]string
 	workers               int
 	kubeConfig            string
 	kubeContext           string
@@ -73,7 +74,7 @@ type kubeControllerCommand struct {
 
 // NewKubeControllerCommand returns the Kubernetes controller command.
 func NewKubeControllerCommand(app *kingpin.Application) Command {
-	c := &kubeControllerCommand{extraLabels: map[string]string{}}
+	c := &kubeControllerCommand{extraLabels: map[string]string{}, idLabels: map[string]string{}}
 	cmd := app.Command("kubernetes-controller", "Runs Sloth in Kubernetes controller/operator mode.")
 	cmd.Alias("controller")
 	cmd.Alias("k8s-controller")
@@ -92,6 +93,7 @@ func NewKubeControllerCommand(app *kingpin.Application) Command {
 	cmd.Flag("hot-reload-addr", "The listen address for hot-reloading components that allow it.").Default(":8082").StringVar(&c.hotReloadAddr)
 	cmd.Flag("hot-reload-path", "The webhook path for hot-reloading components that allow it.").Default("/-/reload").StringVar(&c.hotReloadPath)
 	cmd.Flag("extra-labels", "Extra labels that will be added to all the generated Prometheus rules ('key=value' form, can be repeated).").Short('l').StringMapVar(&c.extraLabels)
+	cmd.Flag("id-labels", "Id labels that used as filters for generated recording rules. These will also be added as extra labels ('key=value' form, can be repeated).").Short('d').StringMapVar(&c.idLabels)
 	cmd.Flag("sli-plugins-path", "The path to SLI plugins (can be repeated), if not set it disable plugins support.").Short('p').StringsVar(&c.sliPluginsPaths)
 	cmd.Flag("slo-period-windows-path", "The directory path to custom SLO period windows catalog (replaces default ones).").StringVar(&c.sloPeriodWindowsPath)
 	cmd.Flag("default-slo-period", "The default SLO period windows to be used for the SLOs.").Default("30d").StringVar(&c.sloPeriod)
@@ -104,6 +106,11 @@ func (k kubeControllerCommand) Name() string { return "kubernetes-controller" }
 func (k kubeControllerCommand) Run(ctx context.Context, config RootConfig) error {
 	logger := config.Logger.WithValues(log.Kv{"window": k.sloPeriod})
 
+	// Make sure id labels are set in extra labels as well
+	for key, value := range k.idLabels {
+		k.extraLabels[key] = value
+	}
+
 	// SLO period.
 	sp, err := prometheusmodel.ParseDuration(k.sloPeriod)
 	if err != nil {
@@ -323,6 +330,7 @@ func (k kubeControllerCommand) Run(ctx context.Context, config RootConfig) error
 			Repository:       k8sprometheus.NewPrometheusOperatorCRDRepo(ksvc, logger),
 			KubeStatusStorer: ksvc,
 			ExtraLabels:      k.extraLabels,
+			IDLabels:         k.idLabels,
 			Logger:           logger,
 		}
 		handler, err := kubecontroller.NewHandler(config)

cmd/sloth/commands/validate.go

@@ -24,19 +24,21 @@ type validateCommand struct {
 	slosExcludeRegex     string
 	slosIncludeRegex     string
 	extraLabels          map[string]string
+	idLabels             map[string]string
 	sliPluginsPaths      []string
 	sloPeriodWindowsPath string
 	sloPeriod            string
 }
 
 // NewValidateCommand returns the validate command.
 func NewValidateCommand(app *kingpin.Application) Command {
-	c := &validateCommand{extraLabels: map[string]string{}}
+	c := &validateCommand{extraLabels: map[string]string{}, idLabels: map[string]string{}}
 	cmd := app.Command("validate", "Validates the SLO manifests and generation of Prometheus SLOs.")
 	cmd.Flag("input", "SLO spec discovery path, will discover recursively all YAML files.").Short('i').Required().StringVar(&c.slosInput)
 	cmd.Flag("fs-exclude", "Filter regex to ignore matched discovered SLO file paths.").Short('e').StringVar(&c.slosExcludeRegex)
 	cmd.Flag("fs-include", "Filter regex to include matched discovered SLO file paths, everything else will be ignored. Exclude has preference.").Short('n').StringVar(&c.slosIncludeRegex)
 	cmd.Flag("extra-labels", "Extra labels that will be added to all the generated Prometheus rules ('key=value' form, can be repeated).").Short('l').StringMapVar(&c.extraLabels)
+	cmd.Flag("id-labels", "Id labels that used as filters for generated recording rules. These will also be added as extra labels ('key=value' form, can be repeated).").Short('d').StringMapVar(&c.idLabels)
 	cmd.Flag("sli-plugins-path", "The path to SLI plugins (can be repeated), if not set it disable plugins support.").Short('p').StringsVar(&c.sliPluginsPaths)
 	cmd.Flag("slo-period-windows-path", "The directory path to custom SLO period windows catalog (replaces default ones).").StringVar(&c.sloPeriodWindowsPath)
 	cmd.Flag("default-slo-period", "The default SLO period windows to be used for the SLOs.").Default("30d").StringVar(&c.sloPeriod)
@@ -48,6 +50,11 @@ func (v validateCommand) Name() string { return "validate" }
 func (v validateCommand) Run(ctx context.Context, config RootConfig) error {
 	logger := config.Logger.WithValues(log.Kv{"window": v.sloPeriod})
 
+	// Make sure id labels are set in extra labels as well
+	for key, value := range v.idLabels {
+		v.extraLabels[key] = value
+	}
+
 	// SLO period.
 	sp, err := prometheusmodel.ParseDuration(v.sloPeriod)
 	if err != nil {
@@ -129,6 +136,7 @@ func (v validateCommand) Run(ctx context.Context, config RootConfig) error {
 			logger:      log.Noop,
 			windowsRepo: windowsRepo,
 			extraLabels: v.extraLabels,
+			idLabels:    v.idLabels,
 		}
 
 		// Prepare file validation result and start validation result for every SLO in the file.

internal/app/generate/prometheus.go

@@ -96,6 +96,8 @@ type Request struct {
 	Info info.Info
 	// ExtraLabels are the extra labels added to the SLOs on execution time.
 	ExtraLabels map[string]string
+	// IDLabels are the extra labels added to the SLOs recording rules on execution time.
+	IDLabels map[string]string
 	// SLOGroup are the SLOs group that will be used to generate the SLO results and Prom rules.
 	SLOGroup prometheus.SLOGroup
 }
@@ -121,6 +123,7 @@ func (s Service) Generate(ctx context.Context, r Request) (*Response, error) {
 	for _, slo := range r.SLOGroup.SLOs {
 		// Add extra labels.
 		slo.Labels = mergeLabels(slo.Labels, r.ExtraLabels)
+		slo.IDLabels = r.IDLabels
 
 		// Generate SLO result.
 		result, err := s.generateSLO(ctx, r.Info, slo)

internal/app/kubecontroller/handler.go

@@ -42,6 +42,7 @@ type HandlerConfig struct {
 	Repository       Repository
 	KubeStatusStorer KubeStatusStorer
 	ExtraLabels      map[string]string
+	IDLabels         map[string]string
 	// IgnoreHandleBefore makes the handles of objects with a success state and no spec change,
 	// be ignored if the last success is less than this setting.
 	// Be aware that this setting should be less than the controller resync interval.
@@ -66,6 +67,10 @@ func (c *HandlerConfig) defaults() error {
 		c.ExtraLabels = map[string]string{}
 	}
 
+	if c.IDLabels == nil {
+		c.IDLabels = map[string]string{}
+	}
+
 	if c.Repository == nil {
 		return fmt.Errorf("repository is required")
 	}
@@ -88,6 +93,7 @@ type handler struct {
 	repository         Repository
 	kubeStatusStorer   KubeStatusStorer
 	extraLabels        map[string]string
+	IDLabels           map[string]string
 	ignoreHandleBefore time.Duration
 	logger             log.Logger
 }
@@ -103,6 +109,7 @@ func NewHandler(config HandlerConfig) (controller.Handler, error) {
 		repository:         config.Repository,
 		kubeStatusStorer:   config.KubeStatusStorer,
 		extraLabels:        config.ExtraLabels,
+		IDLabels:           config.IDLabels,
 		ignoreHandleBefore: config.IgnoreHandleBefore,
 		logger:             config.Logger,
 	}, nil
@@ -152,6 +159,7 @@ func (h handler) handlePrometheusServiceLevelV1(ctx context.Context, psl *slothv
 			Spec:    fmt.Sprintf("%s/%s", slothv1.SchemeGroupVersion.Group, slothv1.SchemeGroupVersion.Version),
 		},
 		ExtraLabels: h.extraLabels,
+		IDLabels:    h.IDLabels,
 		SLOGroup:    model.SLOGroup,
 	}
 	resp, err := h.generator.Generate(ctx, req)

internal/prometheus/alert_rules.go

@@ -101,7 +101,7 @@ func defaultSLOAlertGenerator(slo SLO, sloAlert AlertMeta, quick, slow alert.MWM
 		Alert:       sloAlert.Name,
 		Expr:        expr.String(),
 		Annotations: mergeLabels(extraAnnotations, sloAlert.Annotations),
-		Labels:      mergeLabels(extraLabels, sloAlert.Labels),
+		Labels:      mergeLabels(extraLabels, sloAlert.Labels, slo.IDLabels),
 	}, nil
 }
 

internal/prometheus/model.go

@@ -47,6 +47,7 @@ type SLO struct {
 	TimeWindow      time.Duration     `validate:"required"`
 	Objective       float64           `validate:"gt=0,lte=100"`
 	Labels          map[string]string `validate:"dive,keys,prom_label_key,endkeys,required,prom_label_value"`
+	IDLabels        map[string]string `validate:"dive,keys,prom_label_key,endkeys,required,prom_label_value"`
 	PageAlertMeta   AlertMeta
 	TicketAlertMeta AlertMeta
 }
@@ -68,11 +69,11 @@ func (s SLO) GetSLIErrorMetric(window time.Duration) string {
 // GetSLOIDPromLabels returns the ID labels of an SLO, these can be used to identify
 // an SLO recorded metrics and alerts.
 func (s SLO) GetSLOIDPromLabels() map[string]string {
-	return map[string]string{
+	return mergeLabels(map[string]string{
 		sloIDLabelName:      s.ID,
 		sloNameLabelName:    s.Name,
 		sloServiceLabelName: s.Service,
-	}
+	}, s.IDLabels)
 }
 
 var modelSpecValidate = func() *validator.Validate {