rows.go 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334
  1. package core
  2. import (
  3. "database/sql"
  4. "errors"
  5. "reflect"
  6. "sync"
  7. )
  8. type Rows struct {
  9. *sql.Rows
  10. db *DB
  11. }
  12. func (rs *Rows) ToMapString() ([]map[string]string, error) {
  13. cols, err := rs.Columns()
  14. if err != nil {
  15. return nil, err
  16. }
  17. var results = make([]map[string]string, 0, 10)
  18. for rs.Next() {
  19. var record = make(map[string]string, len(cols))
  20. err = rs.ScanMap(&record)
  21. if err != nil {
  22. return nil, err
  23. }
  24. results = append(results, record)
  25. }
  26. return results, nil
  27. }
  28. // scan data to a struct's pointer according field index
  29. func (rs *Rows) ScanStructByIndex(dest ...interface{}) error {
  30. if len(dest) == 0 {
  31. return errors.New("at least one struct")
  32. }
  33. vvvs := make([]reflect.Value, len(dest))
  34. for i, s := range dest {
  35. vv := reflect.ValueOf(s)
  36. if vv.Kind() != reflect.Ptr || vv.Elem().Kind() != reflect.Struct {
  37. return errors.New("dest should be a struct's pointer")
  38. }
  39. vvvs[i] = vv.Elem()
  40. }
  41. cols, err := rs.Columns()
  42. if err != nil {
  43. return err
  44. }
  45. newDest := make([]interface{}, len(cols))
  46. var i = 0
  47. for _, vvv := range vvvs {
  48. for j := 0; j < vvv.NumField(); j++ {
  49. newDest[i] = vvv.Field(j).Addr().Interface()
  50. i = i + 1
  51. }
  52. }
  53. return rs.Rows.Scan(newDest...)
  54. }
  55. var (
  56. fieldCache = make(map[reflect.Type]map[string]int)
  57. fieldCacheMutex sync.RWMutex
  58. )
  59. func fieldByName(v reflect.Value, name string) reflect.Value {
  60. t := v.Type()
  61. fieldCacheMutex.RLock()
  62. cache, ok := fieldCache[t]
  63. fieldCacheMutex.RUnlock()
  64. if !ok {
  65. cache = make(map[string]int)
  66. for i := 0; i < v.NumField(); i++ {
  67. cache[t.Field(i).Name] = i
  68. }
  69. fieldCacheMutex.Lock()
  70. fieldCache[t] = cache
  71. fieldCacheMutex.Unlock()
  72. }
  73. if i, ok := cache[name]; ok {
  74. return v.Field(i)
  75. }
  76. return reflect.Zero(t)
  77. }
  78. // scan data to a struct's pointer according field name
  79. func (rs *Rows) ScanStructByName(dest interface{}) error {
  80. vv := reflect.ValueOf(dest)
  81. if vv.Kind() != reflect.Ptr || vv.Elem().Kind() != reflect.Struct {
  82. return errors.New("dest should be a struct's pointer")
  83. }
  84. cols, err := rs.Columns()
  85. if err != nil {
  86. return err
  87. }
  88. newDest := make([]interface{}, len(cols))
  89. var v EmptyScanner
  90. for j, name := range cols {
  91. f := fieldByName(vv.Elem(), rs.db.Mapper.Table2Obj(name))
  92. if f.IsValid() {
  93. newDest[j] = f.Addr().Interface()
  94. } else {
  95. newDest[j] = &v
  96. }
  97. }
  98. return rs.Rows.Scan(newDest...)
  99. }
  100. // scan data to a slice's pointer, slice's length should equal to columns' number
  101. func (rs *Rows) ScanSlice(dest interface{}) error {
  102. vv := reflect.ValueOf(dest)
  103. if vv.Kind() != reflect.Ptr || vv.Elem().Kind() != reflect.Slice {
  104. return errors.New("dest should be a slice's pointer")
  105. }
  106. vvv := vv.Elem()
  107. cols, err := rs.Columns()
  108. if err != nil {
  109. return err
  110. }
  111. newDest := make([]interface{}, len(cols))
  112. for j := 0; j < len(cols); j++ {
  113. if j >= vvv.Len() {
  114. newDest[j] = reflect.New(vvv.Type().Elem()).Interface()
  115. } else {
  116. newDest[j] = vvv.Index(j).Addr().Interface()
  117. }
  118. }
  119. err = rs.Rows.Scan(newDest...)
  120. if err != nil {
  121. return err
  122. }
  123. srcLen := vvv.Len()
  124. for i := srcLen; i < len(cols); i++ {
  125. vvv = reflect.Append(vvv, reflect.ValueOf(newDest[i]).Elem())
  126. }
  127. return nil
  128. }
  129. // scan data to a map's pointer
  130. func (rs *Rows) ScanMap(dest interface{}) error {
  131. vv := reflect.ValueOf(dest)
  132. if vv.Kind() != reflect.Ptr || vv.Elem().Kind() != reflect.Map {
  133. return errors.New("dest should be a map's pointer")
  134. }
  135. cols, err := rs.Columns()
  136. if err != nil {
  137. return err
  138. }
  139. newDest := make([]interface{}, len(cols))
  140. vvv := vv.Elem()
  141. for i, _ := range cols {
  142. newDest[i] = rs.db.reflectNew(vvv.Type().Elem()).Interface()
  143. }
  144. err = rs.Rows.Scan(newDest...)
  145. if err != nil {
  146. return err
  147. }
  148. for i, name := range cols {
  149. vname := reflect.ValueOf(name)
  150. vvv.SetMapIndex(vname, reflect.ValueOf(newDest[i]).Elem())
  151. }
  152. return nil
  153. }
  154. type Row struct {
  155. rows *Rows
  156. // One of these two will be non-nil:
  157. err error // deferred error for easy chaining
  158. }
  159. // ErrorRow return an error row
  160. func ErrorRow(err error) *Row {
  161. return &Row{
  162. err: err,
  163. }
  164. }
  165. // NewRow from rows
  166. func NewRow(rows *Rows, err error) *Row {
  167. return &Row{rows, err}
  168. }
  169. func (row *Row) Columns() ([]string, error) {
  170. if row.err != nil {
  171. return nil, row.err
  172. }
  173. return row.rows.Columns()
  174. }
  175. func (row *Row) Scan(dest ...interface{}) error {
  176. if row.err != nil {
  177. return row.err
  178. }
  179. defer row.rows.Close()
  180. for _, dp := range dest {
  181. if _, ok := dp.(*sql.RawBytes); ok {
  182. return errors.New("sql: RawBytes isn't allowed on Row.Scan")
  183. }
  184. }
  185. if !row.rows.Next() {
  186. if err := row.rows.Err(); err != nil {
  187. return err
  188. }
  189. return sql.ErrNoRows
  190. }
  191. err := row.rows.Scan(dest...)
  192. if err != nil {
  193. return err
  194. }
  195. // Make sure the query can be processed to completion with no errors.
  196. return row.rows.Close()
  197. }
  198. func (row *Row) ScanStructByName(dest interface{}) error {
  199. if row.err != nil {
  200. return row.err
  201. }
  202. defer row.rows.Close()
  203. if !row.rows.Next() {
  204. if err := row.rows.Err(); err != nil {
  205. return err
  206. }
  207. return sql.ErrNoRows
  208. }
  209. err := row.rows.ScanStructByName(dest)
  210. if err != nil {
  211. return err
  212. }
  213. // Make sure the query can be processed to completion with no errors.
  214. return row.rows.Close()
  215. }
  216. func (row *Row) ScanStructByIndex(dest interface{}) error {
  217. if row.err != nil {
  218. return row.err
  219. }
  220. defer row.rows.Close()
  221. if !row.rows.Next() {
  222. if err := row.rows.Err(); err != nil {
  223. return err
  224. }
  225. return sql.ErrNoRows
  226. }
  227. err := row.rows.ScanStructByIndex(dest)
  228. if err != nil {
  229. return err
  230. }
  231. // Make sure the query can be processed to completion with no errors.
  232. return row.rows.Close()
  233. }
  234. // scan data to a slice's pointer, slice's length should equal to columns' number
  235. func (row *Row) ScanSlice(dest interface{}) error {
  236. if row.err != nil {
  237. return row.err
  238. }
  239. defer row.rows.Close()
  240. if !row.rows.Next() {
  241. if err := row.rows.Err(); err != nil {
  242. return err
  243. }
  244. return sql.ErrNoRows
  245. }
  246. err := row.rows.ScanSlice(dest)
  247. if err != nil {
  248. return err
  249. }
  250. // Make sure the query can be processed to completion with no errors.
  251. return row.rows.Close()
  252. }
  253. // scan data to a map's pointer
  254. func (row *Row) ScanMap(dest interface{}) error {
  255. if row.err != nil {
  256. return row.err
  257. }
  258. defer row.rows.Close()
  259. if !row.rows.Next() {
  260. if err := row.rows.Err(); err != nil {
  261. return err
  262. }
  263. return sql.ErrNoRows
  264. }
  265. err := row.rows.ScanMap(dest)
  266. if err != nil {
  267. return err
  268. }
  269. // Make sure the query can be processed to completion with no errors.
  270. return row.rows.Close()
  271. }
  272. func (row *Row) ToMapString() (map[string]string, error) {
  273. cols, err := row.Columns()
  274. if err != nil {
  275. return nil, err
  276. }
  277. var record = make(map[string]string, len(cols))
  278. err = row.ScanMap(&record)
  279. if err != nil {
  280. return nil, err
  281. }
  282. return record, nil
  283. }