[B2B-7314] Postgres client for migrations

Author: madatsci

clients/pg.go

@@ -3,11 +3,15 @@ package clients
 import (
 	"context"
 	"crypto/tls"
+	"database/sql"
 	"errors"
+	"fmt"
+	"net/url"
 	"strings"
 	"time"
 
 	"github.com/go-pg/pg/v10"
+	_ "github.com/lib/pq"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"
 )
@@ -18,6 +22,16 @@ type (
 	}
 
 	ctxKey int
+
+	dbConfig struct {
+		addr     string
+		user     string
+		password string
+		database string
+		poolSize int
+		ssl      bool
+		debug    bool
+	}
 )
 
 const ctxRequestStartKey ctxKey = 1 + iota
@@ -47,25 +61,21 @@ func (d dbQueryHook) AfterQuery(ctx context.Context, event *pg.QueryEvent) error
 }
 
 func NewPostgreSQL(config *viper.Viper, logger *logrus.Logger) (*pg.DB, error) {
-	config.SetDefault("database.pool", 10)
-	config.SetDefault("database.debug", false)
-	config.SetDefault("database.ssl", false)
-
-	dbAddr := config.GetString("database.addr")
-	if dbAddr == "" {
-		return nil, errors.New("missing database address")
+	cfg, err := parseDBConfig(config)
+	if err != nil {
+		return nil, err
 	}
 
 	opts := &pg.Options{
-		Addr:     config.GetString("database.addr"),
-		User:     config.GetString("database.user"),
-		Password: config.GetString("database.password"),
-		Database: config.GetString("database.database"),
-		PoolSize: config.GetInt("database.pool"),
+		Addr:     cfg.addr,
+		User:     cfg.user,
+		Password: cfg.password,
+		Database: cfg.database,
+		PoolSize: cfg.poolSize,
 	}
 
-	if config.GetBool("database.ssl") {
-		hp := strings.Split(dbAddr, ":")
+	if cfg.ssl {
+		hp := strings.Split(cfg.addr, ":")
 		if len(hp) != 2 {
 			return nil, errors.New("database address has wrong format")
 		}
@@ -78,7 +88,7 @@ func NewPostgreSQL(config *viper.Viper, logger *logrus.Logger) (*pg.DB, error) {
 
 	connection := pg.Connect(opts)
 
-	if config.GetBool("database.debug") {
+	if cfg.debug {
 		entry := logger.WithField("module", "db")
 		connection.AddQueryHook(dbQueryHook{
 			logger: entry,
@@ -87,3 +97,49 @@ func NewPostgreSQL(config *viper.Viper, logger *logrus.Logger) (*pg.DB, error) {
 
 	return connection, nil
 }
+
+// NewPostgreSQLForMigrations is a connection that is used for migrations.
+// Migrations are implemented with `goose`, which supports only `*sql.DB`.
+func NewPostgreSQLForMigrations(config *viper.Viper) (*sql.DB, error) {
+	cfg, err := parseDBConfig(config)
+	if err != nil {
+		return nil, err
+	}
+
+	dsn := fmt.Sprintf(
+		"postgres://%s:%s@%s/%s",
+		cfg.user,
+		strings.ReplaceAll(url.QueryEscape(cfg.password), ":", "%3A"),
+		cfg.addr,
+		cfg.database,
+	)
+
+	if cfg.ssl {
+		dsn += "?sslmode=verify-ca"
+	} else {
+		dsn += "?sslmode=disable"
+	}
+
+	return sql.Open("postgres", dsn)
+}
+
+func parseDBConfig(config *viper.Viper) (dbConfig, error) {
+	config.SetDefault("database.pool", 10)
+	config.SetDefault("database.debug", false)
+	config.SetDefault("database.ssl", false)
+
+	dbAddr := config.GetString("database.addr")
+	if dbAddr == "" {
+		return dbConfig{}, errors.New("missing database address")
+	}
+
+	return dbConfig{
+		addr:     dbAddr,
+		user:     config.GetString("database.user"),
+		password: config.GetString("database.password"),
+		database: config.GetString("database.database"),
+		poolSize: config.GetInt("database.pool"),
+		ssl:      config.GetBool("database.ssl"),
+		debug:    config.GetBool("database.debug"),
+	}, nil
+}

clients/pg_test.go

@@ -13,13 +13,7 @@ func TestNewPostgreSQL(t *testing.T) {
 		t.Skip("skipping test in short mode")
 	}
 
-	cfg := viper.New()
-	cfg.Set("database.addr", os.Getenv("DATABASE_ADDR"))
-	cfg.Set("database.user", os.Getenv("DATABASE_USER"))
-	cfg.Set("database.password", os.Getenv("DATABASE_PASSWORD"))
-	cfg.Set("database.database", os.Getenv("DATABASE_DATABASE"))
-	cfg.Set("database.ssl", os.Getenv("DATABASE_SSL"))
-
+	cfg := setupConfig()
 	logger := logrus.New()
 
 	db, err := NewPostgreSQL(cfg, logger)
@@ -40,3 +34,40 @@ func TestNewPostgreSQL(t *testing.T) {
 		t.Error("unexpected message")
 	}
 }
+
+func TestNewPostgreSQLForMigrations(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping test in short mode")
+	}
+
+	cfg := setupConfig()
+
+	db, err := NewPostgreSQLForMigrations(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	type StringResult struct {
+		Message string
+	}
+	var res StringResult
+	err = db.QueryRow("SELECT 'hello' AS message").Scan(&res.Message)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if res.Message != "hello" {
+		t.Error("unexpected message")
+	}
+}
+
+func setupConfig() *viper.Viper {
+	cfg := viper.New()
+	cfg.Set("database.addr", os.Getenv("DATABASE_ADDR"))
+	cfg.Set("database.user", os.Getenv("DATABASE_USER"))
+	cfg.Set("database.password", os.Getenv("DATABASE_PASSWORD"))
+	cfg.Set("database.database", os.Getenv("DATABASE_DATABASE"))
+	cfg.Set("database.ssl", os.Getenv("DATABASE_SSL"))
+
+	return cfg
+}

narada/commands/migrations.go

@@ -1,11 +1,10 @@
 package commands
 
 import (
-	"database/sql"
 	"errors"
-	"fmt"
 
 	"github.com/cryptopay-dev/narada"
+	"github.com/cryptopay-dev/narada/clients"
 
 	"github.com/pressly/goose"
 	"github.com/sirupsen/logrus"
@@ -31,7 +30,7 @@ func MigrateUp(p *narada.Narada) *cli.Command {
 				logger.Println("starting migrations")
 				dir := c.String("dir")
 
-				db, err := connect(v)
+				db, err := clients.NewPostgreSQLForMigrations(v)
 				if err != nil {
 					return err
 				}
@@ -61,7 +60,7 @@ func MigrateDown(p *narada.Narada) *cli.Command {
 				logger.Println("rolling back migration")
 				dir := c.String("dir")
 
-				db, err := connect(v)
+				db, err := clients.NewPostgreSQLForMigrations(v)
 				if err != nil {
 					return err
 				}
@@ -107,15 +106,3 @@ func CreateMigration(p *narada.Narada) *cli.Command {
 		},
 	}
 }
-
-func connect(v *viper.Viper) (*sql.DB, error) {
-	dsn := fmt.Sprintf(
-		"postgres://%s:%s@%s/%s?sslmode=disable",
-		v.GetString("database.user"),
-		v.GetString("database.password"),
-		v.GetString("database.addr"),
-		v.GetString("database.database"),
-	)
-
-	return sql.Open("postgres", dsn)
-}