hermez-node/db/utils.go

package db

import (
	"context"
	"database/sql"
	"fmt"
	"math/big"
	"reflect"
	"strings"
	"time"

	"github.com/gobuffalo/packr/v2"
	"github.com/hermeznetwork/hermez-node/log"
	"github.com/hermeznetwork/tracerr"
	"github.com/jmoiron/sqlx"
	migrate "github.com/rubenv/sql-migrate"
	"github.com/russross/meddler"
	"golang.org/x/sync/semaphore"
)

var migrations *migrate.PackrMigrationSource

func init() {
	migrations = &migrate.PackrMigrationSource{
		Box: packr.New("hermez-db-migrations", "./migrations"),
	}
	ms, err := migrations.FindMigrations()
	if err != nil {
		panic(err)
	}
	if len(ms) == 0 {
		panic(fmt.Errorf("no SQL migrations found"))
	}
}

// MigrationsUp runs the SQL migrations Up
func MigrationsUp(db *sql.DB) error {
	nMigrations, err := migrate.Exec(db, "postgres", migrations, migrate.Up)
	if err != nil {
		return tracerr.Wrap(err)
	}
	log.Info("successfully ran ", nMigrations, " migrations Up")
	return nil
}

// MigrationsDown runs the SQL migrations Down
func MigrationsDown(db *sql.DB) error {
	nMigrations, err := migrate.Exec(db, "postgres", migrations, migrate.Down)
	if err != nil {
		return tracerr.Wrap(err)
	}
	log.Info("successfully ran ", nMigrations, " migrations Down")
	return nil
}

// ConnectSQLDB connects to the SQL DB
func ConnectSQLDB(port int, host, user, password, name string) (*sqlx.DB, error) {
	// Init meddler
	initMeddler()
	meddler.Default = meddler.PostgreSQL
	// Stablish connection
	psqlconn := fmt.Sprintf(
		"host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
		host,
		port,
		user,
		password,
		name,
	)
	db, err := sqlx.Connect("postgres", psqlconn)
	if err != nil {
		return nil, tracerr.Wrap(err)
	}
	return db, nil
}

// InitSQLDB runs migrations and registers meddlers
func InitSQLDB(port int, host, user, password, name string) (*sqlx.DB, error) {
	db, err := ConnectSQLDB(port, host, user, password, name)
	if err != nil {
		return nil, tracerr.Wrap(err)
	}
	// Run DB migrations
	if err := MigrationsUp(db.DB); err != nil {
		return nil, tracerr.Wrap(err)
	}
	return db, nil
}

// APIConnectionController is used to limit the SQL open connections used by the API
type APIConnectionController struct {
	smphr   *semaphore.Weighted
	timeout time.Duration
}

// NewAPIConnectionController initialize APIConnectionController
func NewAPIConnectionController(maxConnections int, timeout time.Duration) *APIConnectionController {
	return &APIConnectionController{
		smphr:   semaphore.NewWeighted(int64(maxConnections)),
		timeout: timeout,
	}
}

// Acquire reserves a SQL connection. If the connection is not acquired
// within the timeout, the function will return an error
func (acc *APIConnectionController) Acquire() (context.CancelFunc, error) {
	ctx, cancel := context.WithTimeout(context.Background(), acc.timeout) //nolint:govet
	return cancel, acc.smphr.Acquire(ctx, 1)
}

// Release frees a SQL connection
func (acc *APIConnectionController) Release() {
	acc.smphr.Release(1)
}

// initMeddler registers tags to be used to read/write from SQL DBs using meddler
func initMeddler() {
	meddler.Register("bigint", BigIntMeddler{})
	meddler.Register("bigintnull", BigIntNullMeddler{})
}

// BulkInsert performs a bulk insert with a single statement into the specified table.  Example:
// `db.BulkInsert(myDB, "INSERT INTO block (eth_block_num, timestamp, hash) VALUES %s", blocks[:])`
// Note that all the columns must be specified in the query, and they must be
// in the same order as in the table.
// Note that the fields in the structs need to be defined in the same order as
// in the table columns.
func BulkInsert(db meddler.DB, q string, args interface{}) error {
	arrayValue := reflect.ValueOf(args)
	arrayLen := arrayValue.Len()
	valueStrings := make([]string, 0, arrayLen)
	var arglist = make([]interface{}, 0)
	for i := 0; i < arrayLen; i++ {
		arg := arrayValue.Index(i).Addr().Interface()
		elemArglist, err := meddler.Default.Values(arg, true)
		if err != nil {
			return tracerr.Wrap(err)
		}
		arglist = append(arglist, elemArglist...)
		value := "("
		for j := 0; j < len(elemArglist); j++ {
			value += fmt.Sprintf("$%d, ", i*len(elemArglist)+j+1)
		}
		value = value[:len(value)-2] + ")"
		valueStrings = append(valueStrings, value)
	}
	stmt := fmt.Sprintf(q, strings.Join(valueStrings, ","))
	_, err := db.Exec(stmt, arglist...)
	return tracerr.Wrap(err)
}

// BigIntMeddler encodes or decodes the field value to or from JSON
type BigIntMeddler struct{}

// PreRead is called before a Scan operation for fields that have the BigIntMeddler
func (b BigIntMeddler) PreRead(fieldAddr interface{}) (scanTarget interface{}, err error) {
	// give a pointer to a byte buffer to grab the raw data
	return new(string), nil
}

// PostRead is called after a Scan operation for fields that have the BigIntMeddler
func (b BigIntMeddler) PostRead(fieldPtr, scanTarget interface{}) error {
	ptr := scanTarget.(*string)
	if ptr == nil {
		return tracerr.Wrap(fmt.Errorf("BigIntMeddler.PostRead: nil pointer"))
	}
	field := fieldPtr.(**big.Int)
	*field = new(big.Int).SetBytes([]byte(*ptr))
	return nil
}

// PreWrite is called before an Insert or Update operation for fields that have the BigIntMeddler
func (b BigIntMeddler) PreWrite(fieldPtr interface{}) (saveValue interface{}, err error) {
	field := fieldPtr.(*big.Int)

	return field.Bytes(), nil
}

// BigIntNullMeddler encodes or decodes the field value to or from JSON
type BigIntNullMeddler struct{}

// PreRead is called before a Scan operation for fields that have the BigIntNullMeddler
func (b BigIntNullMeddler) PreRead(fieldAddr interface{}) (scanTarget interface{}, err error) {
	return &fieldAddr, nil
}

// PostRead is called after a Scan operation for fields that have the BigIntNullMeddler
func (b BigIntNullMeddler) PostRead(fieldPtr, scanTarget interface{}) error {
	field := fieldPtr.(**big.Int)
	ptrPtr := scanTarget.(*interface{})
	if *ptrPtr == nil {
		// null column, so set target to be zero value
		*field = nil
		return nil
	}
	// not null
	ptr := (*ptrPtr).([]byte)
	if ptr == nil {
		return tracerr.Wrap(fmt.Errorf("BigIntMeddler.PostRead: nil pointer"))
	}
	*field = new(big.Int).SetBytes(ptr)
	return nil
}

// PreWrite is called before an Insert or Update operation for fields that have the BigIntNullMeddler
func (b BigIntNullMeddler) PreWrite(fieldPtr interface{}) (saveValue interface{}, err error) {
	field := fieldPtr.(*big.Int)
	if field == nil {
		return nil, nil
	}
	return field.Bytes(), nil
}

// SliceToSlicePtrs converts any []Foo to []*Foo
func SliceToSlicePtrs(slice interface{}) interface{} {
	v := reflect.ValueOf(slice)
	vLen := v.Len()
	typ := v.Type().Elem()
	res := reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(typ)), vLen, vLen)
	for i := 0; i < vLen; i++ {
		res.Index(i).Set(v.Index(i).Addr())
	}
	return res.Interface()
}

// SlicePtrsToSlice converts any []*Foo to []Foo
func SlicePtrsToSlice(slice interface{}) interface{} {
	v := reflect.ValueOf(slice)
	vLen := v.Len()
	typ := v.Type().Elem().Elem()
	res := reflect.MakeSlice(reflect.SliceOf(typ), vLen, vLen)
	for i := 0; i < vLen; i++ {
		res.Index(i).Set(v.Index(i).Elem())
	}
	return res.Interface()
}

// Rollback an sql transaction, and log the error if it's not nil
func Rollback(txn *sqlx.Tx) {
	if err := txn.Rollback(); err != nil {
		log.Errorw("Rollback", "err", err)
	}
}

// RowsClose close the rows of an sql query, and log the errir if it's not nil
func RowsClose(rows *sql.Rows) {
	if err := rows.Close(); err != nil {
		log.Errorw("rows.Close", "err", err)
	}
}