path: root/db/db.go

package db

import (
	"fmt"
	"strings"

	"github.com/go-pg/pg"
	"github.com/go-pg/pg/orm"
	"github.com/jloup/utils"
)

var DB *pg.DB

var log = utils.StandardL().WithField("module", "db")

type DBConfig struct {
	Address  string
	Database string
	User     string
	Password string
}

func Init(config DBConfig) {
	var err error

	DB = connect(config)

	err = createSchema(DB)
	if err != nil {
		log.Errorf("cannot create schemas %v\n", err)
	}

	err = createIndexes(DB)
	if err != nil {
		log.Errorf("cannot create indexes %v\n", err)
	}
}

func connect(config DBConfig) *pg.DB {
	return pg.Connect(&pg.Options{
		User:     config.User,
		Password: config.Password,
		Database: config.Database,
		Addr:     config.Address,
	})
}

func createSchema(db *pg.DB) error {
	for _, model := range []interface{}{&User{}, &MarketConfig{}} {
		err := db.CreateTable(model, &orm.CreateTableOptions{IfNotExists: true})
		if err != nil {
			return err
		}
	}
	return nil
}

func createIndexes(db *pg.DB) error {
	indexes := []struct {
		TableName string
		Name      string
		Columns   []string
	}{
		{"market_configs", "market_name_user_id_idx", []string{"user_id", "market_name"}},
	}

	for _, index := range indexes {
		_, err := db.Exec(fmt.Sprintf("CREATE UNIQUE INDEX IF NOT EXISTS %s ON %s (%s)", index.Name, index.TableName, strings.Join(index.Columns, ",")))
		if err != nil {
			return err
		}
	}

	return nil
}