arnaucube
arnaucube/hermez-nodePublic Self-hosted repository
hermez-node/db/utils.go
db/utils.goRaw
  1package db
  2
  3import (
  4	"context"
  5	"database/sql"
  6	"fmt"
  7	"math/big"
  8	"reflect"
  9	"strings"
 10	"time"
 11
 12	"github.com/gobuffalo/packr/v2"
 13	"github.com/hermeznetwork/hermez-node/log"
 14	"github.com/hermeznetwork/tracerr"
 15	"github.com/jmoiron/sqlx"
 16	migrate "github.com/rubenv/sql-migrate"
 17	"github.com/russross/meddler"
 18	"golang.org/x/sync/semaphore"
 19)
 20
 21var migrations *migrate.PackrMigrationSource
 22
 23func init() {
 24	migrations = &migrate.PackrMigrationSource{
 25		Box: packr.New("hermez-db-migrations", "./migrations"),
 26	}
 27	ms, err := migrations.FindMigrations()
 28	if err != nil {
 29		panic(err)
 30	}
 31	if len(ms) == 0 {
 32		panic(fmt.Errorf("no SQL migrations found"))
 33	}
 34}
 35
 36// MigrationsUp runs the SQL migrations Up
 37func MigrationsUp(db *sql.DB) error {
 38	nMigrations, err := migrate.Exec(db, "postgres", migrations, migrate.Up)
 39	if err != nil {
 40		return tracerr.Wrap(err)
 41	}
 42	log.Info("successfully ran ", nMigrations, " migrations Up")
 43	return nil
 44}
 45
 46// MigrationsDown runs the SQL migrations Down
 47func MigrationsDown(db *sql.DB) error {
 48	nMigrations, err := migrate.Exec(db, "postgres", migrations, migrate.Down)
 49	if err != nil {
 50		return tracerr.Wrap(err)
 51	}
 52	log.Info("successfully ran ", nMigrations, " migrations Down")
 53	return nil
 54}
 55
 56// ConnectSQLDB connects to the SQL DB
 57func ConnectSQLDB(port int, host, user, password, name string) (*sqlx.DB, error) {
 58	// Init meddler
 59	initMeddler()
 60	meddler.Default = meddler.PostgreSQL
 61	// Stablish connection
 62	psqlconn := fmt.Sprintf(
 63		"host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
 64		host,
 65		port,
 66		user,
 67		password,
 68		name,
 69	)
 70	db, err := sqlx.Connect("postgres", psqlconn)
 71	if err != nil {
 72		return nil, tracerr.Wrap(err)
 73	}
 74	return db, nil
 75}
 76
 77// InitSQLDB runs migrations and registers meddlers
 78func InitSQLDB(port int, host, user, password, name string) (*sqlx.DB, error) {
 79	db, err := ConnectSQLDB(port, host, user, password, name)
 80	if err != nil {
 81		return nil, tracerr.Wrap(err)
 82	}
 83	// Run DB migrations
 84	if err := MigrationsUp(db.DB); err != nil {
 85		return nil, tracerr.Wrap(err)
 86	}
 87	return db, nil
 88}
 89
 90// APIConnectionController is used to limit the SQL open connections used by the API
 91type APIConnectionController struct {
 92	smphr   *semaphore.Weighted
 93	timeout time.Duration
 94}
 95
 96// NewAPICnnectionController initialize APIConnectionController
 97func NewAPICnnectionController(maxConnections int, timeout time.Duration) *APIConnectionController {
 98	return &APIConnectionController{
 99		smphr:   semaphore.NewWeighted(int64(maxConnections)),
100		timeout: timeout,
101	}
102}
103
104// Acquire reserves a SQL connection. If the connection is not acquired
105// within the timeout, the function will return an error
106func (acc *APIConnectionController) Acquire() (context.CancelFunc, error) {
107	ctx, cancel := context.WithTimeout(context.Background(), acc.timeout) //nolint:govet
108	return cancel, acc.smphr.Acquire(ctx, 1)
109}
110
111// Release frees a SQL connection
112func (acc *APIConnectionController) Release() {
113	acc.smphr.Release(1)
114}
115
116// initMeddler registers tags to be used to read/write from SQL DBs using meddler
117func initMeddler() {
118	meddler.Register("bigint", BigIntMeddler{})
119	meddler.Register("bigintnull", BigIntNullMeddler{})
120}
121
122// BulkInsert performs a bulk insert with a single statement into the specified table.  Example:
123// `db.BulkInsert(myDB, "INSERT INTO block (eth_block_num, timestamp, hash) VALUES %s", blocks[:])`
124// Note that all the columns must be specified in the query, and they must be
125// in the same order as in the table.
126// Note that the fields in the structs need to be defined in the same order as
127// in the table columns.
128func BulkInsert(db meddler.DB, q string, args interface{}) error {
129	arrayValue := reflect.ValueOf(args)
130	arrayLen := arrayValue.Len()
131	valueStrings := make([]string, 0, arrayLen)
132	var arglist = make([]interface{}, 0)
133	for i := 0; i < arrayLen; i++ {
134		arg := arrayValue.Index(i).Addr().Interface()
135		elemArglist, err := meddler.Default.Values(arg, true)
136		if err != nil {
137			return tracerr.Wrap(err)
138		}
139		arglist = append(arglist, elemArglist...)
140		value := "("
141		for j := 0; j < len(elemArglist); j++ {
142			value += fmt.Sprintf("$%d, ", i*len(elemArglist)+j+1)
143		}
144		value = value[:len(value)-2] + ")"
145		valueStrings = append(valueStrings, value)
146	}
147	stmt := fmt.Sprintf(q, strings.Join(valueStrings, ","))
148	_, err := db.Exec(stmt, arglist...)
149	return tracerr.Wrap(err)
150}
151
152// BigIntMeddler encodes or decodes the field value to or from JSON
153type BigIntMeddler struct{}
154
155// PreRead is called before a Scan operation for fields that have the BigIntMeddler
156func (b BigIntMeddler) PreRead(fieldAddr interface{}) (scanTarget interface{}, err error) {
157	// give a pointer to a byte buffer to grab the raw data
158	return new(string), nil
159}
160
161// PostRead is called after a Scan operation for fields that have the BigIntMeddler
162func (b BigIntMeddler) PostRead(fieldPtr, scanTarget interface{}) error {
163	ptr := scanTarget.(*string)
164	if ptr == nil {
165		return tracerr.Wrap(fmt.Errorf("BigIntMeddler.PostRead: nil pointer"))
166	}
167	field := fieldPtr.(**big.Int)
168	*field = new(big.Int).SetBytes([]byte(*ptr))
169	return nil
170}
171
172// PreWrite is called before an Insert or Update operation for fields that have the BigIntMeddler
173func (b BigIntMeddler) PreWrite(fieldPtr interface{}) (saveValue interface{}, err error) {
174	field := fieldPtr.(*big.Int)
175
176	return field.Bytes(), nil
177}
178
179// BigIntNullMeddler encodes or decodes the field value to or from JSON
180type BigIntNullMeddler struct{}
181
182// PreRead is called before a Scan operation for fields that have the BigIntNullMeddler
183func (b BigIntNullMeddler) PreRead(fieldAddr interface{}) (scanTarget interface{}, err error) {
184	return &fieldAddr, nil
185}
186
187// PostRead is called after a Scan operation for fields that have the BigIntNullMeddler
188func (b BigIntNullMeddler) PostRead(fieldPtr, scanTarget interface{}) error {
189	field := fieldPtr.(**big.Int)
190	ptrPtr := scanTarget.(*interface{})
191	if *ptrPtr == nil {
192		// null column, so set target to be zero value
193		*field = nil
194		return nil
195	}
196	// not null
197	ptr := (*ptrPtr).([]byte)
198	if ptr == nil {
199		return tracerr.Wrap(fmt.Errorf("BigIntMeddler.PostRead: nil pointer"))
200	}
201	*field = new(big.Int).SetBytes(ptr)
202	return nil
203}
204
205// PreWrite is called before an Insert or Update operation for fields that have the BigIntNullMeddler
206func (b BigIntNullMeddler) PreWrite(fieldPtr interface{}) (saveValue interface{}, err error) {
207	field := fieldPtr.(*big.Int)
208	if field == nil {
209		return nil, nil
210	}
211	return field.Bytes(), nil
212}
213
214// SliceToSlicePtrs converts any []Foo to []*Foo
215func SliceToSlicePtrs(slice interface{}) interface{} {
216	v := reflect.ValueOf(slice)
217	vLen := v.Len()
218	typ := v.Type().Elem()
219	res := reflect.MakeSlice(reflect.SliceOf(reflect.PtrTo(typ)), vLen, vLen)
220	for i := 0; i < vLen; i++ {
221		res.Index(i).Set(v.Index(i).Addr())
222	}
223	return res.Interface()
224}
225
226// SlicePtrsToSlice converts any []*Foo to []Foo
227func SlicePtrsToSlice(slice interface{}) interface{} {
228	v := reflect.ValueOf(slice)
229	vLen := v.Len()
230	typ := v.Type().Elem().Elem()
231	res := reflect.MakeSlice(reflect.SliceOf(typ), vLen, vLen)
232	for i := 0; i < vLen; i++ {
233		res.Index(i).Set(v.Index(i).Elem())
234	}
235	return res.Interface()
236}
237
238// Rollback an sql transaction, and log the error if it's not nil
239func Rollback(txn *sqlx.Tx) {
240	if err := txn.Rollback(); err != nil {
241		log.Errorw("Rollback", "err", err)
242	}
243}
244
245// RowsClose close the rows of an sql query, and log the errir if it's not nil
246func RowsClose(rows *sql.Rows) {
247	if err := rows.Close(); err != nil {
248		log.Errorw("rows.Close", "err", err)
249	}
250}