]>
Commit | Line | Data |
---|---|---|
1 | package db | |
2 | ||
3 | import ( | |
4 | "fmt" | |
5 | "strings" | |
6 | ||
7 | "github.com/go-pg/pg" | |
8 | "github.com/go-pg/pg/orm" | |
9 | "github.com/jloup/utils" | |
10 | ) | |
11 | ||
12 | var DB *pg.DB | |
13 | ||
14 | var log = utils.StandardL().WithField("module", "db") | |
15 | ||
16 | type DBConfig struct { | |
17 | Address string | |
18 | Database string | |
19 | User string | |
20 | Password string | |
21 | } | |
22 | ||
23 | func Init(config DBConfig) { | |
24 | var err error | |
25 | ||
26 | DB = connect(config) | |
27 | ||
28 | err = createSchema(DB) | |
29 | if err != nil { | |
30 | log.Errorf("cannot create schemas %v\n", err) | |
31 | } | |
32 | ||
33 | err = createIndexes(DB) | |
34 | if err != nil { | |
35 | log.Errorf("cannot create indexes %v\n", err) | |
36 | } | |
37 | } | |
38 | ||
39 | func connect(config DBConfig) *pg.DB { | |
40 | return pg.Connect(&pg.Options{ | |
41 | User: config.User, | |
42 | Password: config.Password, | |
43 | Database: config.Database, | |
44 | Addr: config.Address, | |
45 | }) | |
46 | } | |
47 | ||
48 | func createSchema(db *pg.DB) error { | |
49 | for _, model := range []interface{}{&User{}, &MarketConfig{}} { | |
50 | err := db.CreateTable(model, &orm.CreateTableOptions{IfNotExists: true}) | |
51 | if err != nil { | |
52 | return err | |
53 | } | |
54 | } | |
55 | return nil | |
56 | } | |
57 | ||
58 | func createIndexes(db *pg.DB) error { | |
59 | indexes := []struct { | |
60 | TableName string | |
61 | Name string | |
62 | Columns []string | |
63 | }{ | |
64 | {"market_configs", "market_name_user_id_idx", []string{"user_id", "market_name"}}, | |
65 | } | |
66 | ||
67 | for _, index := range indexes { | |
68 | _, err := db.Exec(fmt.Sprintf("CREATE UNIQUE INDEX IF NOT EXISTS %s ON %s (%s)", index.Name, index.TableName, strings.Join(index.Columns, ","))) | |
69 | if err != nil { | |
70 | return err | |
71 | } | |
72 | } | |
73 | ||
74 | return nil | |
75 | } |