db/db.go

   1 package db
   2
   3 import (
   4         migrate "github.com/go-pg/migrations"
   5         "github.com/go-pg/pg"
   6         "github.com/go-pg/pg/orm"
   7         "github.com/go-redis/redis"
   8         "github.com/jloup/utils"
   9 )
  10
  11 var log = utils.StandardL().WithField("module", "db")
  12
  13 var (
  14         DB    *pg.DB
  15         Redis *redis.Client
  16 )
  17
  18 type DBConfig struct {
  19         Type     string // tcp or unix
  20         Address  string
  21         Database string
  22         User     string
  23         Password string
  24 }
  25
  26 type RedisConfig struct {
  27         Type     string // tcp or unix
  28         Address  string
  29         Password string
  30         Database int
  31 }
  32
  33 func Init(config DBConfig, redisConfig RedisConfig) {
  34         var err error
  35
  36         DB = connect(config)
  37
  38         err = migratedb()
  39         if err != nil {
  40                 log.Fatalf("cannot migratedb '%v'\n", err)
  41         }
  42
  43         Redis = redis.NewClient(&redis.Options{
  44                 Network:  redisConfig.Type,
  45                 Addr:     redisConfig.Address,
  46                 Password: redisConfig.Password,
  47                 DB:       redisConfig.Database,
  48         })
  49
  50         _, err = Redis.Ping().Result()
  51
  52         if err != nil {
  53                 log.Fatalf("redis init error %s", err)
  54         }
  55
  56 }
  57
  58 func migratedb() error {
  59
  60         mig := make([]migrate.Migration, 0)
  61
  62         for i := range migrations {
  63                 index := i
  64                 mig = append(mig, migrate.Migration{
  65                         Version: migrations[index].Version,
  66                         Up: func(db orm.DB) error {
  67                                 for _, query := range migrations[index].Up {
  68                                         _, err := db.Exec(query)
  69                                         if err != nil {
  70                                                 return err
  71                                         }
  72                                 }
  73
  74                                 return nil
  75                         },
  76                         Down: func(db orm.DB) error {
  77                                 for _, query := range migrations[index].Down {
  78                                         _, err := db.Exec(query)
  79                                         if err != nil {
  80                                                 return err
  81                                         }
  82                                 }
  83
  84                                 return nil
  85                         },
  86                 })
  87         }
  88
  89         oldVersion, newVersion, err := migrate.RunMigrations(DB, mig, "up")
  90
  91         if oldVersion != newVersion {
  92                 log.Infof("Migrate DB: %v -> %v", oldVersion, newVersion)
  93         } else {
  94                 log.Infof("DB up-to-date: version '%v'", newVersion)
  95         }
  96         return err
  97 }
  98
  99 func connect(config DBConfig) *pg.DB {
 100         return pg.Connect(&pg.Options{
 101                 User:     config.User,
 102                 Password: config.Password,
 103                 Database: config.Database,
 104                 Network:  config.Type,
 105                 Addr:     config.Address,
 106         })
 107 }