Separate DB interface and implementation for rules configs

Author: jml

pkg/configs/configs.go

@@ -7,8 +7,8 @@ type ID int
 // A Config is a Cortex configuration for a single user.
 type Config struct {
 	// RulesFiles maps from a rules filename to file contents.
-	RulesFiles         map[string]string `json:"rules_files"`
-	AlertmanagerConfig string            `json:"alertmanager_config"`
+	RulesFiles         RulesConfig `json:"rules_files"`
+	AlertmanagerConfig string      `json:"alertmanager_config"`
 }
 
 // View is what's returned from the Weave Cloud configs service
@@ -21,3 +21,21 @@ type View struct {
 	ID     ID     `json:"id"`
 	Config Config `json:"config"`
 }
+
+// GetVersionedRulesConfig specializes the view to just the rules config.
+func (v View) GetVersionedRulesConfig() VersionedRulesConfig {
+	return VersionedRulesConfig{
+		ID:     v.ID,
+		Config: v.Config.RulesFiles,
+	}
+}
+
+// RulesConfig are the set of rules files for a particular organization.
+type RulesConfig map[string]string
+
+// VersionedRulesConfig is a RulesConfig together with a version.
+// `data Versioned a = Versioned { id :: ID , config :: a }`
+type VersionedRulesConfig struct {
+	ID     ID          `json:"id"`
+	Config RulesConfig `json:"config"`
+}

pkg/configs/db/db.go

@@ -22,8 +22,24 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	flag.StringVar(&cfg.MigrationsDir, "database.migrations", "", "Path where the database migration files can be found")
 }
 
+// RulesDB has ruler-specific DB interfaces.
+type RulesDB interface {
+	// GetRulesConfig gets the user's ruler config
+	GetRulesConfig(userID string) (configs.VersionedRulesConfig, error)
+	SetRulesConfig(userID string, config configs.RulesConfig) error
+	// SetRulesConfig sets the user's ruler config
+
+	// GetAllRulesConfigs gets all of the ruler configs
+	GetAllRulesConfigs() (map[string]configs.VersionedRulesConfig, error)
+	// GetRulesConfigs gets all of the configs that have been added or have
+	// changed since the provided config.
+	GetRulesConfigs(since configs.ID) (map[string]configs.VersionedRulesConfig, error)
+}
+
 // DB is the interface for the database.
 type DB interface {
+	RulesDB
+
 	GetConfig(userID string) (configs.View, error)
 	SetConfig(userID string, cfg configs.Config) error
 

pkg/configs/db/memory/memory.go

@@ -56,3 +56,46 @@ func (d *DB) GetConfigs(since configs.ID) (map[string]configs.View, error) {
 func (d *DB) Close() error {
 	return nil
 }
+
+// GetRulesConfig gets the rules config for a user.
+func (d *DB) GetRulesConfig(userID string) (configs.VersionedRulesConfig, error) {
+	c, ok := d.cfgs[userID]
+	if !ok {
+		return configs.VersionedRulesConfig{}, sql.ErrNoRows
+	}
+	return c.GetVersionedRulesConfig(), nil
+}
+
+// SetRulesConfig sets the rules config for a user.
+func (d *DB) SetRulesConfig(userID string, config configs.RulesConfig) error {
+	c, ok := d.cfgs[userID]
+	if !ok {
+		return d.SetConfig(userID, configs.Config{RulesFiles: config})
+	}
+	return d.SetConfig(userID, configs.Config{
+		AlertmanagerConfig: c.Config.AlertmanagerConfig,
+		RulesFiles:         config,
+	})
+}
+
+// GetAllRulesConfigs gets the rules configs for all users that have them.
+func (d *DB) GetAllRulesConfigs() (map[string]configs.VersionedRulesConfig, error) {
+	cfgs := map[string]configs.VersionedRulesConfig{}
+	for user, c := range d.cfgs {
+		cfgs[user] = c.GetVersionedRulesConfig()
+	}
+	return cfgs, nil
+}
+
+// GetRulesConfigs gets the rules configs that have changed
+// since the given config version.
+func (d *DB) GetRulesConfigs(since configs.ID) (map[string]configs.VersionedRulesConfig, error) {
+	cfgs := map[string]configs.VersionedRulesConfig{}
+	for user, c := range d.cfgs {
+		if c.ID <= since {
+			continue
+		}
+		cfgs[user] = c.GetVersionedRulesConfig()
+	}
+	return cfgs, nil
+}

pkg/configs/db/postgres/postgres.go

@@ -134,6 +134,71 @@ func (d DB) GetConfigs(since configs.ID) (map[string]configs.View, error) {
 	})
 }
 
+// GetRulesConfig gets the latest alertmanager config for a user.
+func (d DB) GetRulesConfig(userID string) (configs.VersionedRulesConfig, error) {
+	current, err := d.GetConfig(userID)
+	if err != nil {
+		return configs.VersionedRulesConfig{}, err
+	}
+	return current.GetVersionedRulesConfig(), nil
+}
+
+// SetRulesConfig sets the current alertmanager config for a user.
+func (d DB) SetRulesConfig(userID string, config configs.RulesConfig) error {
+	current, err := d.GetConfig(userID)
+	if err != nil {
+		return err
+	}
+	new := configs.Config{
+		AlertmanagerConfig: current.Config.AlertmanagerConfig,
+		RulesFiles:         config,
+	}
+	return d.SetConfig(userID, new)
+}
+
+func (d DB) findRulesConfigs(filter squirrel.Sqlizer) (map[string]configs.VersionedRulesConfig, error) {
+	rows, err := d.Select("id", "owner_id", "config ->> 'rules_files'").
+		Options("DISTINCT ON (owner_id)").
+		From("configs").
+		Where(filter).
+		Where("config ->> 'rules_files' <> '{}'").
+		OrderBy("owner_id, id DESC").
+		Query()
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	cfgs := map[string]configs.VersionedRulesConfig{}
+	for rows.Next() {
+		var cfg configs.VersionedRulesConfig
+		var userID string
+		var cfgBytes []byte
+		err = rows.Scan(&cfg.ID, &userID, &cfgBytes)
+		if err != nil {
+			return nil, err
+		}
+		err = json.Unmarshal(cfgBytes, &cfg.Config)
+		if err != nil {
+			return nil, err
+		}
+		cfgs[userID] = cfg
+	}
+	return cfgs, nil
+}
+
+// GetAllRulesConfigs gets all alertmanager configs for all users.
+func (d DB) GetAllRulesConfigs() (map[string]configs.VersionedRulesConfig, error) {
+	return d.findRulesConfigs(activeConfig)
+}
+
+// GetRulesConfigs gets all the alertmanager configs that have changed since a given config.
+func (d DB) GetRulesConfigs(since configs.ID) (map[string]configs.VersionedRulesConfig, error) {
+	return d.findRulesConfigs(squirrel.And{
+		activeConfig,
+		squirrel.Gt{"id": since},
+	})
+}
+
 // Transaction runs the given function in a postgres transaction. If fn returns
 // an error the txn will be rolled back.
 func (d DB) Transaction(f func(DB) error) error {

pkg/configs/db/timed.go

@@ -74,3 +74,33 @@ func (t timed) Close() error {
 		return t.d.Close()
 	})
 }
+
+func (t timed) GetRulesConfig(userID string) (cfg configs.VersionedRulesConfig, err error) {
+	t.timeRequest("GetRulesConfig", func(_ context.Context) error {
+		cfg, err = t.d.GetRulesConfig(userID)
+		return err
+	})
+	return
+}
+
+func (t timed) SetRulesConfig(userID string, cfg configs.RulesConfig) (err error) {
+	return t.timeRequest("SetRulesConfig", func(_ context.Context) error {
+		return t.d.SetRulesConfig(userID, cfg)
+	})
+}
+
+func (t timed) GetAllRulesConfigs() (cfgs map[string]configs.VersionedRulesConfig, err error) {
+	t.timeRequest("GetAllRulesConfigs", func(_ context.Context) error {
+		cfgs, err = t.d.GetAllRulesConfigs()
+		return err
+	})
+	return
+}
+
+func (t timed) GetRulesConfigs(since configs.ID) (cfgs map[string]configs.VersionedRulesConfig, err error) {
+	t.timeRequest("GetRulesConfigs", func(_ context.Context) error {
+		cfgs, err = t.d.GetRulesConfigs(since)
+		return err
+	})
+	return
+}

pkg/configs/db/traced.go

@@ -41,3 +41,23 @@ func (t traced) Close() (err error) {
 	defer func() { t.trace("Close", err) }()
 	return t.d.Close()
 }
+
+func (t traced) GetRulesConfig(userID string) (cfg configs.VersionedRulesConfig, err error) {
+	defer func() { t.trace("GetRulesConfig", userID, cfg, err) }()
+	return t.d.GetRulesConfig(userID)
+}
+
+func (t traced) SetRulesConfig(userID string, cfg configs.RulesConfig) (err error) {
+	defer func() { t.trace("SetRulesConfig", userID, cfg, err) }()
+	return t.d.SetRulesConfig(userID, cfg)
+}
+
+func (t traced) GetAllRulesConfigs() (cfgs map[string]configs.VersionedRulesConfig, err error) {
+	defer func() { t.trace("GetAllRulesConfigs", cfgs, err) }()
+	return t.d.GetAllRulesConfigs()
+}
+
+func (t traced) GetRulesConfigs(since configs.ID) (cfgs map[string]configs.VersionedRulesConfig, err error) {
+	defer func() { t.trace("GetConfigs", since, cfgs, err) }()
+	return t.d.GetRulesConfigs(since)
+}