You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

250 lines
7.3 KiB

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"
"github.com/marusama/semaphore/v2"
migrate "github.com/rubenv/sql-migrate"
"github.com/russross/meddler"
)
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.Semaphore
timeout time.Duration
}
// NewAPICnnectionController initialize APIConnectionController
func NewAPICnnectionController(maxConnections int, timeout time.Duration) *APIConnectionController {
return &APIConnectionController{
smphr: semaphore.New(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)
}
}