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.

175 lines
5.3 KiB

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