111 lines
2.2 KiB
Go
111 lines
2.2 KiB
Go
package db
|
|
|
|
import (
|
|
"fmt"
|
|
"sync"
|
|
|
|
"fiskerinc.com/modules/logger"
|
|
"fiskerinc.com/modules/utils/envtool"
|
|
"github.com/go-pg/pg/v10"
|
|
"github.com/go-pg/pg/v10/orm"
|
|
"github.com/pkg/errors"
|
|
|
|
pgtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/go-pg/pg.v10"
|
|
)
|
|
|
|
var (
|
|
ErrTableDoesntExist = errors.New("table doesn't exist")
|
|
ErrTableAndModelDontMatch = errors.New("table and model columns don't match")
|
|
)
|
|
|
|
func GetDefaultConn() *pg.DB {
|
|
host := envtool.GetEnv("DB_HOST", "localhost")
|
|
port := envtool.GetEnv("DB_PORT", "5432")
|
|
user := envtool.GetEnv("DB_USER", "postgres")
|
|
password := envtool.GetEnv("DB_PASSWORD", "REPLACE_ME")
|
|
dbname := envtool.GetEnv("DB_NAME", "postgres")
|
|
sslmode := envtool.GetEnv("DB_SSLMODE", "disable")
|
|
poolSize := envtool.GetEnvInt("DB_POOLSIZE", 10)
|
|
addr := fmt.Sprintf("%v:%v", host, port)
|
|
|
|
logger.Info().Msgf("Initializing database connection %s", addr)
|
|
conn_str := fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=%s", user, password, host, port, dbname, sslmode)
|
|
opts, err := pg.ParseURL(conn_str)
|
|
opts.PoolSize = poolSize
|
|
if err != nil {
|
|
logger.Error().Err(err).Send()
|
|
}
|
|
conn := pg.Connect(opts)
|
|
|
|
// Wrap the connection with the APM hook.
|
|
pgtrace.Wrap(conn)
|
|
|
|
return conn
|
|
}
|
|
|
|
type DBClientInterface interface {
|
|
GetConn() *pg.DB
|
|
SetConn(*pg.DB) error
|
|
InitSchema([]interface{}) error
|
|
RegisterManyToManyRel(tables []interface{})
|
|
Close() error
|
|
}
|
|
|
|
type DBClient struct {
|
|
driver *pg.DB
|
|
once sync.Once
|
|
}
|
|
|
|
func (db *DBClient) Close() error {
|
|
logger.Info().Msg("Closing database connection")
|
|
|
|
if db.driver == nil {
|
|
return nil
|
|
}
|
|
|
|
return db.driver.Close()
|
|
}
|
|
|
|
func (db *DBClient) GetConn() *pg.DB {
|
|
db.once.Do(func() {
|
|
if db.driver != nil {
|
|
return
|
|
}
|
|
db.driver = GetDefaultConn()
|
|
})
|
|
|
|
return db.driver
|
|
}
|
|
|
|
func (db *DBClient) SetConn(d *pg.DB) error {
|
|
var err error
|
|
if db.driver != nil {
|
|
err = db.driver.Close()
|
|
}
|
|
|
|
db.driver = d
|
|
return err
|
|
}
|
|
|
|
func (db *DBClient) InitSchema(models []interface{}) error {
|
|
migrator := Migrator{
|
|
DB: db.GetConn(),
|
|
}
|
|
defer migrator.Close()
|
|
|
|
for _, model := range models {
|
|
err := migrator.Check(model)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (db *DBClient) RegisterManyToManyRel(tables []interface{}) {
|
|
for _, table := range tables {
|
|
orm.RegisterTable(table)
|
|
}
|
|
}
|