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

     1  package db
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"net/url"
     9  
    10  	"github.com/cortexproject/cortex/pkg/configs/db/memory"
    11  	"github.com/cortexproject/cortex/pkg/configs/db/postgres"
    12  	"github.com/cortexproject/cortex/pkg/configs/userconfig"
    13  )
    14  
    15  // Config configures the database.
    16  type Config struct {
    17  	URI           string `yaml:"uri"`
    18  	MigrationsDir string `yaml:"migrations_dir"`
    19  	PasswordFile  string `yaml:"password_file"`
    20  
    21  	// Allow injection of mock DBs for unit testing.
    22  	Mock DB `yaml:"-"`
    23  }
    24  
    25  // RegisterFlags adds the flags required to configure this to the given FlagSet.
    26  func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
    27  	f.StringVar(&cfg.URI, "configs.database.uri", "postgres://postgres@configs-db.weave.local/configs?sslmode=disable", "URI where the database can be found (for dev you can use memory://)")
    28  	f.StringVar(&cfg.MigrationsDir, "configs.database.migrations-dir", "", "Path where the database migration files can be found")
    29  	f.StringVar(&cfg.PasswordFile, "configs.database.password-file", "", "File containing password (username goes in URI)")
    30  }
    31  
    32  // DB is the interface for the database.
    33  type DB interface {
    34  	// GetRulesConfig gets the user's ruler config
    35  	GetRulesConfig(ctx context.Context, userID string) (userconfig.VersionedRulesConfig, error)
    36  
    37  	// SetRulesConfig does a compare-and-swap (CAS) on the user's rules config.
    38  	// `oldConfig` must precisely match the current config in order to change the config to `newConfig`.
    39  	// Will return `true` if the config was updated, `false` otherwise.
    40  	SetRulesConfig(ctx context.Context, userID string, oldConfig, newConfig userconfig.RulesConfig) (bool, error)
    41  
    42  	// GetAllRulesConfigs gets all of the ruler configs
    43  	GetAllRulesConfigs(ctx context.Context) (map[string]userconfig.VersionedRulesConfig, error)
    44  
    45  	// GetRulesConfigs gets all of the configs that have been added or have
    46  	// changed since the provided config.
    47  	GetRulesConfigs(ctx context.Context, since userconfig.ID) (map[string]userconfig.VersionedRulesConfig, error)
    48  
    49  	GetConfig(ctx context.Context, userID string) (userconfig.View, error)
    50  	SetConfig(ctx context.Context, userID string, cfg userconfig.Config) error
    51  
    52  	GetAllConfigs(ctx context.Context) (map[string]userconfig.View, error)
    53  	GetConfigs(ctx context.Context, since userconfig.ID) (map[string]userconfig.View, error)
    54  
    55  	DeactivateConfig(ctx context.Context, userID string) error
    56  	RestoreConfig(ctx context.Context, userID string) error
    57  
    58  	Close() error
    59  }
    60  
    61  // New creates a new database.
    62  func New(cfg Config) (DB, error) {
    63  	if cfg.Mock != nil {
    64  		return cfg.Mock, nil
    65  	}
    66  
    67  	u, err := url.Parse(cfg.URI)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  
    72  	if len(cfg.PasswordFile) != 0 {
    73  		if u.User == nil {
    74  			return nil, fmt.Errorf("--database.password-file requires username in --database.uri")
    75  		}
    76  		passwordBytes, err := ioutil.ReadFile(cfg.PasswordFile)
    77  		if err != nil {
    78  			return nil, fmt.Errorf("Could not read database password file: %v", err)
    79  		}
    80  		u.User = url.UserPassword(u.User.Username(), string(passwordBytes))
    81  	}
    82  
    83  	var d DB
    84  	switch u.Scheme {
    85  	case "memory":
    86  		d, err = memory.New(u.String(), cfg.MigrationsDir)
    87  	case "postgres":
    88  		d, err = postgres.New(u.String(), cfg.MigrationsDir)
    89  	default:
    90  		return nil, fmt.Errorf("Unknown database type: %s", u.Scheme)
    91  	}
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  	return traced{timed{d}}, nil
    96  }