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