arnaucube
arnaucube/hermez-nodePublic Self-hosted repository
hermez-node/db/utils.go
db/utils.goRaw
  1package db
  2
  3import (
  4	"encoding/base64"
  5	"fmt"
  6	"math/big"
  7	"reflect"
  8	"strings"
  9
 10	"github.com/gobuffalo/packr/v2"
 11	"github.com/hermeznetwork/hermez-node/log"
 12	"github.com/jmoiron/sqlx"
 13	migrate "github.com/rubenv/sql-migrate"
 14	"github.com/russross/meddler"
 15)
 16
 17// InitSQLDB runs migrations and registers meddlers
 18func InitSQLDB(port int, host, user, password, name string) (*sqlx.DB, error) {
 19	// Init meddler
 20	initMeddler()
 21	meddler.Default = meddler.PostgreSQL
 22	// Stablish connection
 23	psqlconn := fmt.Sprintf(
 24		"host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
 25		host,
 26		port,
 27		user,
 28		password,
 29		name,
 30	)
 31	db, err := sqlx.Connect("postgres", psqlconn)
 32	if err != nil {
 33		return nil, err
 34	}
 35	// Run DB migrations
 36	migrations := &migrate.PackrMigrationSource{
 37		Box: packr.New("hermez-db-migrations", "./migrations"),
 38	}
 39	nMigrations, err := migrate.Exec(db.DB, "postgres", migrations, migrate.Up)
 40	if err != nil {
 41		return nil, err
 42	}
 43	log.Info("successfully runt ", nMigrations, " migrations")
 44	return db, nil
 45}
 46
 47// initMeddler registers tags to be used to read/write from SQL DBs using meddler
 48func initMeddler() {
 49	meddler.Register("bigint", BigIntMeddler{})
 50	meddler.Register("bigintnull", BigIntNullMeddler{})
 51}
 52
 53// BulkInsert performs a bulk insert with a single statement into the specified table.  Example:
 54// `db.BulkInsert(myDB, "INSERT INTO block (eth_block_num, timestamp, hash) VALUES %s", blocks[:])`
 55// Note that all the columns must be specified in the query, and they must be in the same order as in the table.
 56func BulkInsert(db meddler.DB, q string, args interface{}) error {
 57	arrayValue := reflect.ValueOf(args)
 58	arrayLen := arrayValue.Len()
 59	valueStrings := make([]string, 0, arrayLen)
 60	var arglist = make([]interface{}, 0)
 61	for i := 0; i < arrayLen; i++ {
 62		arg := arrayValue.Index(i).Addr().Interface()
 63		elemArglist, err := meddler.Default.Values(arg, true)
 64		if err != nil {
 65			return err
 66		}
 67		arglist = append(arglist, elemArglist...)
 68		value := "("
 69		for j := 0; j < len(elemArglist); j++ {
 70			value += fmt.Sprintf("$%d, ", i*len(elemArglist)+j+1)
 71		}
 72		value = value[:len(value)-2] + ")"
 73		valueStrings = append(valueStrings, value)
 74	}
 75	stmt := fmt.Sprintf(q, strings.Join(valueStrings, ","))
 76	_, err := db.Exec(stmt, arglist...)
 77	return err
 78}
 79
 80// BigIntMeddler encodes or decodes the field value to or from JSON
 81type BigIntMeddler struct{}
 82
 83// PreRead is called before a Scan operation for fields that have the BigIntMeddler
 84func (b BigIntMeddler) PreRead(fieldAddr interface{}) (scanTarget interface{}, err error) {
 85	// give a pointer to a byte buffer to grab the raw data
 86	return new(string), nil
 87}
 88
 89// PostRead is called after a Scan operation for fields that have the BigIntMeddler
 90func (b BigIntMeddler) PostRead(fieldPtr, scanTarget interface{}) error {
 91	ptr := scanTarget.(*string)
 92	if ptr == nil {
 93		return fmt.Errorf("BigIntMeddler.PostRead: nil pointer")
 94	}
 95
 96	data, err := base64.StdEncoding.DecodeString(*ptr)
 97	if err != nil {
 98		return fmt.Errorf("big.Int decode error: %v", err)
 99	}
100	field := fieldPtr.(**big.Int)
101	*field = new(big.Int).SetBytes(data)
102
103	return nil
104}
105
106// PreWrite is called before an Insert or Update operation for fields that have the BigIntMeddler
107func (b BigIntMeddler) PreWrite(fieldPtr interface{}) (saveValue interface{}, err error) {
108	field := fieldPtr.(*big.Int)
109
110	str := base64.StdEncoding.EncodeToString(field.Bytes())
111
112	return str, nil
113}
114
115// BigIntNullMeddler encodes or decodes the field value to or from JSON
116type BigIntNullMeddler struct{}
117
118// PreRead is called before a Scan operation for fields that have the BigIntNullMeddler
119func (b BigIntNullMeddler) PreRead(fieldAddr interface{}) (scanTarget interface{}, err error) {
120	return &fieldAddr, nil
121}
122
123// PostRead is called after a Scan operation for fields that have the BigIntNullMeddler
124func (b BigIntNullMeddler) PostRead(fieldPtr, scanTarget interface{}) error {
125	sv := reflect.ValueOf(scanTarget)
126	if sv.Elem().IsNil() {
127		// null column, so set target to be zero value
128		fv := reflect.ValueOf(fieldPtr)
129		fv.Elem().Set(reflect.Zero(fv.Elem().Type()))
130		return nil
131	}
132	// not null
133	encoded := new([]byte)
134	refEnc := reflect.ValueOf(encoded)
135	refEnc.Elem().Set(sv.Elem().Elem())
136	data, err := base64.StdEncoding.DecodeString(string(*encoded))
137	if err != nil {
138		return fmt.Errorf("big.Int decode error: %v", err)
139	}
140	field := fieldPtr.(**big.Int)
141	*field = new(big.Int).SetBytes(data)
142	return nil
143}
144
145// PreWrite is called before an Insert or Update operation for fields that have the BigIntNullMeddler
146func (b BigIntNullMeddler) PreWrite(fieldPtr interface{}) (saveValue interface{}, err error) {
147	field := fieldPtr.(*big.Int)
148	if field == nil {
149		return nil, nil
150	}
151	return base64.StdEncoding.EncodeToString(field.Bytes()), nil
152}