db.go 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227
  1. // Copyright 2019 The Xorm Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package core
  5. import (
  6. "context"
  7. "database/sql"
  8. "database/sql/driver"
  9. "fmt"
  10. "reflect"
  11. "regexp"
  12. "sync"
  13. )
  14. var (
  15. // DefaultCacheSize sets the default cache size
  16. DefaultCacheSize = 200
  17. )
  18. func MapToSlice(query string, mp interface{}) (string, []interface{}, error) {
  19. vv := reflect.ValueOf(mp)
  20. if vv.Kind() != reflect.Ptr || vv.Elem().Kind() != reflect.Map {
  21. return "", []interface{}{}, ErrNoMapPointer
  22. }
  23. args := make([]interface{}, 0, len(vv.Elem().MapKeys()))
  24. var err error
  25. query = re.ReplaceAllStringFunc(query, func(src string) string {
  26. v := vv.Elem().MapIndex(reflect.ValueOf(src[1:]))
  27. if !v.IsValid() {
  28. err = fmt.Errorf("map key %s is missing", src[1:])
  29. } else {
  30. args = append(args, v.Interface())
  31. }
  32. return "?"
  33. })
  34. return query, args, err
  35. }
  36. func StructToSlice(query string, st interface{}) (string, []interface{}, error) {
  37. vv := reflect.ValueOf(st)
  38. if vv.Kind() != reflect.Ptr || vv.Elem().Kind() != reflect.Struct {
  39. return "", []interface{}{}, ErrNoStructPointer
  40. }
  41. args := make([]interface{}, 0)
  42. var err error
  43. query = re.ReplaceAllStringFunc(query, func(src string) string {
  44. fv := vv.Elem().FieldByName(src[1:]).Interface()
  45. if v, ok := fv.(driver.Valuer); ok {
  46. var value driver.Value
  47. value, err = v.Value()
  48. if err != nil {
  49. return "?"
  50. }
  51. args = append(args, value)
  52. } else {
  53. args = append(args, fv)
  54. }
  55. return "?"
  56. })
  57. if err != nil {
  58. return "", []interface{}{}, err
  59. }
  60. return query, args, nil
  61. }
  62. type cacheStruct struct {
  63. value reflect.Value
  64. idx int
  65. }
  66. // DB is a wrap of sql.DB with extra contents
  67. type DB struct {
  68. *sql.DB
  69. Mapper IMapper
  70. reflectCache map[reflect.Type]*cacheStruct
  71. reflectCacheMutex sync.RWMutex
  72. }
  73. // Open opens a database
  74. func Open(driverName, dataSourceName string) (*DB, error) {
  75. db, err := sql.Open(driverName, dataSourceName)
  76. if err != nil {
  77. return nil, err
  78. }
  79. return &DB{
  80. DB: db,
  81. Mapper: NewCacheMapper(&SnakeMapper{}),
  82. reflectCache: make(map[reflect.Type]*cacheStruct),
  83. }, nil
  84. }
  85. // FromDB creates a DB from a sql.DB
  86. func FromDB(db *sql.DB) *DB {
  87. return &DB{
  88. DB: db,
  89. Mapper: NewCacheMapper(&SnakeMapper{}),
  90. reflectCache: make(map[reflect.Type]*cacheStruct),
  91. }
  92. }
  93. func (db *DB) reflectNew(typ reflect.Type) reflect.Value {
  94. db.reflectCacheMutex.Lock()
  95. defer db.reflectCacheMutex.Unlock()
  96. cs, ok := db.reflectCache[typ]
  97. if !ok || cs.idx+1 > DefaultCacheSize-1 {
  98. cs = &cacheStruct{reflect.MakeSlice(reflect.SliceOf(typ), DefaultCacheSize, DefaultCacheSize), 0}
  99. db.reflectCache[typ] = cs
  100. } else {
  101. cs.idx = cs.idx + 1
  102. }
  103. return cs.value.Index(cs.idx).Addr()
  104. }
  105. // QueryContext overwrites sql.DB.QueryContext
  106. func (db *DB) QueryContext(ctx context.Context, query string, args ...interface{}) (*Rows, error) {
  107. rows, err := db.DB.QueryContext(ctx, query, args...)
  108. if err != nil {
  109. if rows != nil {
  110. rows.Close()
  111. }
  112. return nil, err
  113. }
  114. return &Rows{rows, db}, nil
  115. }
  116. // Query overwrites sql.DB.Query
  117. func (db *DB) Query(query string, args ...interface{}) (*Rows, error) {
  118. return db.QueryContext(context.Background(), query, args...)
  119. }
  120. // QueryMapContext executes query with parameters via map and context
  121. func (db *DB) QueryMapContext(ctx context.Context, query string, mp interface{}) (*Rows, error) {
  122. query, args, err := MapToSlice(query, mp)
  123. if err != nil {
  124. return nil, err
  125. }
  126. return db.QueryContext(ctx, query, args...)
  127. }
  128. // QueryMap executes query with parameters via map
  129. func (db *DB) QueryMap(query string, mp interface{}) (*Rows, error) {
  130. return db.QueryMapContext(context.Background(), query, mp)
  131. }
  132. func (db *DB) QueryStructContext(ctx context.Context, query string, st interface{}) (*Rows, error) {
  133. query, args, err := StructToSlice(query, st)
  134. if err != nil {
  135. return nil, err
  136. }
  137. return db.QueryContext(ctx, query, args...)
  138. }
  139. func (db *DB) QueryStruct(query string, st interface{}) (*Rows, error) {
  140. return db.QueryStructContext(context.Background(), query, st)
  141. }
  142. func (db *DB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *Row {
  143. rows, err := db.QueryContext(ctx, query, args...)
  144. if err != nil {
  145. return &Row{nil, err}
  146. }
  147. return &Row{rows, nil}
  148. }
  149. func (db *DB) QueryRow(query string, args ...interface{}) *Row {
  150. return db.QueryRowContext(context.Background(), query, args...)
  151. }
  152. func (db *DB) QueryRowMapContext(ctx context.Context, query string, mp interface{}) *Row {
  153. query, args, err := MapToSlice(query, mp)
  154. if err != nil {
  155. return &Row{nil, err}
  156. }
  157. return db.QueryRowContext(ctx, query, args...)
  158. }
  159. func (db *DB) QueryRowMap(query string, mp interface{}) *Row {
  160. return db.QueryRowMapContext(context.Background(), query, mp)
  161. }
  162. func (db *DB) QueryRowStructContext(ctx context.Context, query string, st interface{}) *Row {
  163. query, args, err := StructToSlice(query, st)
  164. if err != nil {
  165. return &Row{nil, err}
  166. }
  167. return db.QueryRowContext(ctx, query, args...)
  168. }
  169. func (db *DB) QueryRowStruct(query string, st interface{}) *Row {
  170. return db.QueryRowStructContext(context.Background(), query, st)
  171. }
  172. var (
  173. re = regexp.MustCompile(`[?](\w+)`)
  174. )
  175. // ExecMapContext exec map with context.Context
  176. // insert into (name) values (?)
  177. // insert into (name) values (?name)
  178. func (db *DB) ExecMapContext(ctx context.Context, query string, mp interface{}) (sql.Result, error) {
  179. query, args, err := MapToSlice(query, mp)
  180. if err != nil {
  181. return nil, err
  182. }
  183. return db.DB.ExecContext(ctx, query, args...)
  184. }
  185. func (db *DB) ExecMap(query string, mp interface{}) (sql.Result, error) {
  186. return db.ExecMapContext(context.Background(), query, mp)
  187. }
  188. func (db *DB) ExecStructContext(ctx context.Context, query string, st interface{}) (sql.Result, error) {
  189. query, args, err := StructToSlice(query, st)
  190. if err != nil {
  191. return nil, err
  192. }
  193. return db.DB.ExecContext(ctx, query, args...)
  194. }
  195. func (db *DB) ExecStruct(query string, st interface{}) (sql.Result, error) {
  196. return db.ExecStructContext(context.Background(), query, st)
  197. }