github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/configs/db/memory/memory.go (about)

     1  package memory
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/cortexproject/cortex/pkg/configs/userconfig"
    10  )
    11  
    12  // DB is an in-memory database for testing, and local development
    13  type DB struct {
    14  	cfgs map[string]userconfig.View
    15  	id   uint
    16  }
    17  
    18  // New creates a new in-memory database
    19  func New(_, _ string) (*DB, error) {
    20  	return &DB{
    21  		cfgs: map[string]userconfig.View{},
    22  		id:   0,
    23  	}, nil
    24  }
    25  
    26  // GetConfig gets the user's configuration.
    27  func (d *DB) GetConfig(ctx context.Context, userID string) (userconfig.View, error) {
    28  	c, ok := d.cfgs[userID]
    29  	if !ok {
    30  		return userconfig.View{}, sql.ErrNoRows
    31  	}
    32  	return c, nil
    33  }
    34  
    35  // SetConfig sets configuration for a user.
    36  func (d *DB) SetConfig(ctx context.Context, userID string, cfg userconfig.Config) error {
    37  	if !cfg.RulesConfig.FormatVersion.IsValid() {
    38  		return fmt.Errorf("invalid rule format version %v", cfg.RulesConfig.FormatVersion)
    39  	}
    40  	d.cfgs[userID] = userconfig.View{Config: cfg, ID: userconfig.ID(d.id)}
    41  	d.id++
    42  	return nil
    43  }
    44  
    45  // GetAllConfigs gets all of the userconfig.
    46  func (d *DB) GetAllConfigs(ctx context.Context) (map[string]userconfig.View, error) {
    47  	return d.cfgs, nil
    48  }
    49  
    50  // GetConfigs gets all of the configs that have changed recently.
    51  func (d *DB) GetConfigs(ctx context.Context, since userconfig.ID) (map[string]userconfig.View, error) {
    52  	cfgs := map[string]userconfig.View{}
    53  	for user, c := range d.cfgs {
    54  		if c.ID > since {
    55  			cfgs[user] = c
    56  		}
    57  	}
    58  	return cfgs, nil
    59  }
    60  
    61  // SetDeletedAtConfig sets a deletedAt for configuration
    62  // by adding a single new row with deleted_at set
    63  // the same as SetConfig is actually insert
    64  func (d *DB) SetDeletedAtConfig(ctx context.Context, userID string, deletedAt time.Time) error {
    65  	cv, err := d.GetConfig(ctx, userID)
    66  	if err != nil {
    67  		return err
    68  	}
    69  	cv.DeletedAt = deletedAt
    70  	cv.ID = userconfig.ID(d.id)
    71  	d.cfgs[userID] = cv
    72  	d.id++
    73  	return nil
    74  }
    75  
    76  // DeactivateConfig deactivates configuration for a user by creating new configuration with DeletedAt set to now
    77  func (d *DB) DeactivateConfig(ctx context.Context, userID string) error {
    78  	return d.SetDeletedAtConfig(ctx, userID, time.Now())
    79  }
    80  
    81  // RestoreConfig restores deactivated configuration for a user by creating new configuration with empty DeletedAt
    82  func (d *DB) RestoreConfig(ctx context.Context, userID string) error {
    83  	return d.SetDeletedAtConfig(ctx, userID, time.Time{})
    84  }
    85  
    86  // Close finishes using the db. Noop.
    87  func (d *DB) Close() error {
    88  	return nil
    89  }
    90  
    91  // GetRulesConfig gets the rules config for a user.
    92  func (d *DB) GetRulesConfig(ctx context.Context, userID string) (userconfig.VersionedRulesConfig, error) {
    93  	c, ok := d.cfgs[userID]
    94  	if !ok {
    95  		return userconfig.VersionedRulesConfig{}, sql.ErrNoRows
    96  	}
    97  	cfg := c.GetVersionedRulesConfig()
    98  	if cfg == nil {
    99  		return userconfig.VersionedRulesConfig{}, sql.ErrNoRows
   100  	}
   101  	return *cfg, nil
   102  }
   103  
   104  // SetRulesConfig sets the rules config for a user.
   105  func (d *DB) SetRulesConfig(ctx context.Context, userID string, oldConfig, newConfig userconfig.RulesConfig) (bool, error) {
   106  	c, ok := d.cfgs[userID]
   107  	if !ok {
   108  		return true, d.SetConfig(ctx, userID, userconfig.Config{RulesConfig: newConfig})
   109  	}
   110  	if !oldConfig.Equal(c.Config.RulesConfig) {
   111  		return false, nil
   112  	}
   113  	return true, d.SetConfig(ctx, userID, userconfig.Config{
   114  		AlertmanagerConfig: c.Config.AlertmanagerConfig,
   115  		RulesConfig:        newConfig,
   116  	})
   117  }
   118  
   119  // GetAllRulesConfigs gets the rules configs for all users that have them.
   120  func (d *DB) GetAllRulesConfigs(ctx context.Context) (map[string]userconfig.VersionedRulesConfig, error) {
   121  	cfgs := map[string]userconfig.VersionedRulesConfig{}
   122  	for user, c := range d.cfgs {
   123  		cfg := c.GetVersionedRulesConfig()
   124  		if cfg != nil {
   125  			cfgs[user] = *cfg
   126  		}
   127  	}
   128  	return cfgs, nil
   129  }
   130  
   131  // GetRulesConfigs gets the rules configs that have changed
   132  // since the given config version.
   133  func (d *DB) GetRulesConfigs(ctx context.Context, since userconfig.ID) (map[string]userconfig.VersionedRulesConfig, error) {
   134  	cfgs := map[string]userconfig.VersionedRulesConfig{}
   135  	for user, c := range d.cfgs {
   136  		if c.ID <= since {
   137  			continue
   138  		}
   139  		cfg := c.GetVersionedRulesConfig()
   140  		if cfg != nil {
   141  			cfgs[user] = *cfg
   142  		}
   143  	}
   144  	return cfgs, nil
   145  }