Initial cloud-services repo - gateway service + pkg modules

pkg/db/db.go

@@ -0,0 +1,110 @@
+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)
+	}
+}