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.

118 lines
3.8 KiB

  1. package db
  2. import (
  3. "encoding/base64"
  4. "fmt"
  5. "math/big"
  6. "reflect"
  7. "strings"
  8. "github.com/russross/meddler"
  9. )
  10. // InitMeddler registers tags to be used to read/write from SQL DBs using meddler
  11. func InitMeddler() {
  12. meddler.Register("bigint", BigIntMeddler{})
  13. meddler.Register("bigintnull", BigIntNullMeddler{})
  14. }
  15. // BulkInsert performs a bulk insert with a single statement into the specified table. Example:
  16. // `db.BulkInsert(myDB, "INSERT INTO block (eth_block_num, timestamp, hash) VALUES %s", blocks[:])`
  17. // Note that all the columns must be specified in the query, and they must be in the same order as in the table.
  18. func BulkInsert(db meddler.DB, q string, args interface{}) error {
  19. arrayValue := reflect.ValueOf(args)
  20. arrayLen := arrayValue.Len()
  21. valueStrings := make([]string, 0, arrayLen)
  22. var arglist = make([]interface{}, 0)
  23. for i := 0; i < arrayLen; i++ {
  24. arg := arrayValue.Index(i).Addr().Interface()
  25. elemArglist, err := meddler.Default.Values(arg, true)
  26. if err != nil {
  27. return err
  28. }
  29. arglist = append(arglist, elemArglist...)
  30. value := "("
  31. for j := 0; j < len(elemArglist); j++ {
  32. value += fmt.Sprintf("$%d, ", i*len(elemArglist)+j+1)
  33. }
  34. value = value[:len(value)-2] + ")"
  35. valueStrings = append(valueStrings, value)
  36. }
  37. stmt := fmt.Sprintf(q, strings.Join(valueStrings, ","))
  38. _, err := db.Exec(stmt, arglist...)
  39. return err
  40. }
  41. // BigIntMeddler encodes or decodes the field value to or from JSON
  42. type BigIntMeddler struct{}
  43. // PreRead is called before a Scan operation for fields that have the BigIntMeddler
  44. func (b BigIntMeddler) PreRead(fieldAddr interface{}) (scanTarget interface{}, err error) {
  45. // give a pointer to a byte buffer to grab the raw data
  46. return new(string), nil
  47. }
  48. // PostRead is called after a Scan operation for fields that have the BigIntMeddler
  49. func (b BigIntMeddler) PostRead(fieldPtr, scanTarget interface{}) error {
  50. ptr := scanTarget.(*string)
  51. if ptr == nil {
  52. return fmt.Errorf("BigIntMeddler.PostRead: nil pointer")
  53. }
  54. data, err := base64.StdEncoding.DecodeString(*ptr)
  55. if err != nil {
  56. return fmt.Errorf("big.Int decode error: %v", err)
  57. }
  58. field := fieldPtr.(**big.Int)
  59. *field = new(big.Int).SetBytes(data)
  60. return nil
  61. }
  62. // PreWrite is called before an Insert or Update operation for fields that have the BigIntMeddler
  63. func (b BigIntMeddler) PreWrite(fieldPtr interface{}) (saveValue interface{}, err error) {
  64. field := fieldPtr.(*big.Int)
  65. str := base64.StdEncoding.EncodeToString(field.Bytes())
  66. return str, nil
  67. }
  68. // BigIntNullMeddler encodes or decodes the field value to or from JSON
  69. type BigIntNullMeddler struct{}
  70. // PreRead is called before a Scan operation for fields that have the BigIntNullMeddler
  71. func (b BigIntNullMeddler) PreRead(fieldAddr interface{}) (scanTarget interface{}, err error) {
  72. return &fieldAddr, nil
  73. }
  74. // PostRead is called after a Scan operation for fields that have the BigIntNullMeddler
  75. func (b BigIntNullMeddler) PostRead(fieldPtr, scanTarget interface{}) error {
  76. sv := reflect.ValueOf(scanTarget)
  77. if sv.Elem().IsNil() {
  78. // null column, so set target to be zero value
  79. fv := reflect.ValueOf(fieldPtr)
  80. fv.Elem().Set(reflect.Zero(fv.Elem().Type()))
  81. return nil
  82. }
  83. // not null
  84. encoded := new([]byte)
  85. refEnc := reflect.ValueOf(encoded)
  86. refEnc.Elem().Set(sv.Elem().Elem())
  87. data, err := base64.StdEncoding.DecodeString(string(*encoded))
  88. if err != nil {
  89. return fmt.Errorf("big.Int decode error: %v", err)
  90. }
  91. field := fieldPtr.(**big.Int)
  92. *field = new(big.Int).SetBytes(data)
  93. return nil
  94. }
  95. // PreWrite is called before an Insert or Update operation for fields that have the BigIntNullMeddler
  96. func (b BigIntNullMeddler) PreWrite(fieldPtr interface{}) (saveValue interface{}, err error) {
  97. field := fieldPtr.(*big.Int)
  98. if field == nil {
  99. return nil, nil
  100. }
  101. return base64.StdEncoding.EncodeToString(field.Bytes()), nil
  102. }