/
connection.go
82 lines (71 loc) · 2.3 KB
/
connection.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
package sql
import (
"context"
"reflect"
"github.com/caarlos0/env"
"github.com/devtron-labs/lens/internal/logger"
pg "github.com/go-pg/pg/v10"
"go.uber.org/zap"
)
type Config struct {
Addr string `env:"PG_ADDR" envDefault:"127.0.0.1"`
Port string `env:"PG_PORT" envDefault:"5432"`
User string `env:"PG_USER" envDefault:""`
Password string `env:"PG_PASSWORD" envDefault:"" secretData:"-"`
Database string `env:"PG_DATABASE" envDefault:"lens"`
ApplicationName string `env:"APP" envDefault:"lens"`
LogQuery bool `env:"PG_LOG_QUERY" envDefault:"true"`
}
func (d dbLogger) BeforeQuery(c context.Context, q *pg.QueryEvent) (context.Context, error) {
return c, nil
}
func (d dbLogger) AfterQuery(c context.Context, q *pg.QueryEvent) error {
query, err := q.FormattedQuery()
logger.NewSugardLogger().Debugw("Printing formatted query", "query", query)
return err
}
type dbLogger struct {
beforeQueryMethod func(context.Context, *pg.QueryEvent) (context.Context, error)
afterQueryMethod func(context.Context, *pg.QueryEvent) error
}
func GetConfig() (*Config, error) {
cfg := &Config{}
err := env.Parse(cfg)
return cfg, err
}
func NewDbConnection(cfg *Config, logger *zap.SugaredLogger) (*pg.DB, error) {
options := pg.Options{
Addr: cfg.Addr + ":" + cfg.Port,
User: cfg.User,
Password: cfg.Password,
Database: cfg.Database,
ApplicationName: cfg.ApplicationName,
}
dbConnection := pg.Connect(&options)
//check db connection
var test string
_, err := dbConnection.QueryOne(pg.Scan(&test), "SELECT 1")
if err != nil {
logger.Errorw("error in connecting db ", "db", obfuscateSecretTags(cfg), "err", err)
return nil, err
} else {
logger.Infow("connected with db", "db", obfuscateSecretTags(cfg))
}
if cfg.LogQuery {
dbConnection.AddQueryHook(dbLogger{})
}
return dbConnection, err
}
func obfuscateSecretTags(cfg interface{}) interface{} {
cfgDpl := reflect.New(reflect.ValueOf(cfg).Elem().Type()).Interface()
cfgDplElm := reflect.ValueOf(cfgDpl).Elem()
t := cfgDplElm.Type()
for i := 0; i < t.NumField(); i++ {
if _, ok := t.Field(i).Tag.Lookup("secretData"); ok {
cfgDplElm.Field(i).SetString("********")
} else {
cfgDplElm.Field(i).Set(reflect.ValueOf(cfg).Elem().Field(i))
}
}
return cfgDpl
}