mongodb.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865
  1. package mongodb
  2. import (
  3. "context"
  4. "encoding/json"
  5. "fmt"
  6. "log"
  7. "math/big"
  8. "reflect"
  9. "runtime"
  10. "strconv"
  11. "strings"
  12. "time"
  13. "go.mongodb.org/mongo-driver/bson"
  14. "go.mongodb.org/mongo-driver/bson/primitive"
  15. "go.mongodb.org/mongo-driver/mongo"
  16. "go.mongodb.org/mongo-driver/mongo/options"
  17. )
  18. func NewMgo(addr, db string, size int) *MongodbSim {
  19. mgo := &MongodbSim{
  20. MongodbAddr: addr,
  21. Size: size,
  22. DbName: db,
  23. }
  24. mgo.InitPool()
  25. return mgo
  26. }
  27. func NewMgoWithUser(addr, db, uname, upwd string, size int) *MongodbSim {
  28. mgo := &MongodbSim{
  29. MongodbAddr: addr,
  30. Size: size,
  31. DbName: db,
  32. UserName: uname,
  33. Password: upwd,
  34. }
  35. mgo.InitPool()
  36. return mgo
  37. }
  38. type Bluk struct {
  39. ms *MgoSess
  40. writes []mongo.WriteModel
  41. }
  42. func (b *Bluk) Insert(doc interface{}) {
  43. write := mongo.NewInsertOneModel()
  44. write.SetDocument(doc)
  45. b.writes = append(b.writes, write)
  46. }
  47. func (b *Bluk) Update(doc ...interface{}) {
  48. write := mongo.NewUpdateOneModel()
  49. write.SetFilter(doc[0])
  50. ue := ObjToM(doc[1])
  51. autoUpdateTime(b.ms.db, b.ms.coll, ue)
  52. write.SetUpdate(ue)
  53. write.SetUpsert(false)
  54. b.writes = append(b.writes, write)
  55. }
  56. func (b *Bluk) UpdateAll(doc ...interface{}) {
  57. write := mongo.NewUpdateManyModel()
  58. write.SetFilter(doc[0])
  59. ue := ObjToM(doc[1])
  60. autoUpdateTime(b.ms.db, b.ms.coll, ue)
  61. write.SetUpdate(ue)
  62. write.SetUpsert(false)
  63. b.writes = append(b.writes, write)
  64. }
  65. func (b *Bluk) Upsert(doc ...interface{}) {
  66. write := mongo.NewUpdateOneModel()
  67. write.SetFilter(doc[0])
  68. ue := ObjToM(doc[1])
  69. autoUpdateTime(b.ms.db, b.ms.coll, ue)
  70. write.SetUpdate(ue)
  71. write.SetUpsert(true)
  72. b.writes = append(b.writes, write)
  73. }
  74. func (b *Bluk) Remove(doc interface{}) {
  75. write := mongo.NewDeleteOneModel()
  76. write.SetFilter(doc)
  77. b.writes = append(b.writes, write)
  78. }
  79. func (b *Bluk) RemoveAll(doc interface{}) {
  80. write := mongo.NewDeleteManyModel()
  81. write.SetFilter(doc)
  82. b.writes = append(b.writes, write)
  83. }
  84. func (b *Bluk) Run() (*mongo.BulkWriteResult, error) {
  85. return b.ms.M.C.Database(b.ms.db).Collection(b.ms.coll).BulkWrite(b.ms.M.Ctx, b.writes)
  86. }
  87. //
  88. type MgoIter struct {
  89. Cursor *mongo.Cursor
  90. Ctx context.Context
  91. }
  92. func (mt *MgoIter) Next(result interface{}) bool {
  93. if mt.Cursor != nil {
  94. if mt.Cursor.Next(mt.Ctx) {
  95. rType := reflect.TypeOf(result)
  96. rVal := reflect.ValueOf(result)
  97. if rType.Kind() == reflect.Ptr {
  98. rType = rType.Elem()
  99. rVal = rVal.Elem()
  100. }
  101. var err error
  102. if rType.Kind() == reflect.Map {
  103. r := make(map[string]interface{})
  104. err = mt.Cursor.Decode(&r)
  105. if rVal.CanSet() {
  106. rVal.Set(reflect.ValueOf(r))
  107. } else {
  108. for it := rVal.MapRange(); it.Next(); {
  109. rVal.SetMapIndex(it.Key(), reflect.Value{})
  110. }
  111. for it := reflect.ValueOf(r).MapRange(); it.Next(); {
  112. rVal.SetMapIndex(it.Key(), it.Value())
  113. }
  114. }
  115. } else {
  116. err = mt.Cursor.Decode(&result)
  117. }
  118. if err != nil {
  119. log.Println("mgo cur err", err.Error())
  120. mt.Cursor.Close(mt.Ctx)
  121. return false
  122. }
  123. return true
  124. } else {
  125. mt.Cursor.Close(mt.Ctx)
  126. return false
  127. }
  128. } else {
  129. return false
  130. }
  131. }
  132. //
  133. type MgoSess struct {
  134. db string
  135. coll string
  136. query interface{}
  137. sorts []string
  138. fields interface{}
  139. limit int64
  140. skip int64
  141. all interface{}
  142. M *MongodbSim
  143. }
  144. func (ms *MgoSess) DB(name string) *MgoSess {
  145. ms.db = name
  146. return ms
  147. }
  148. func (ms *MgoSess) C(name string) *MgoSess {
  149. ms.coll = name
  150. return ms
  151. }
  152. func (ms *MgoSess) Bulk() *Bluk {
  153. return &Bluk{ms: ms}
  154. }
  155. func (ms *MgoSess) Find(q interface{}) *MgoSess {
  156. if q == nil {
  157. q = map[string]interface{}{}
  158. }
  159. ms.query = q
  160. return ms
  161. }
  162. func (ms *MgoSess) FindId(_id interface{}) *MgoSess {
  163. ms.query = map[string]interface{}{"_id": _id}
  164. return ms
  165. }
  166. func (ms *MgoSess) Select(fields interface{}) *MgoSess {
  167. ms.fields = fields
  168. return ms
  169. }
  170. func (ms *MgoSess) Limit(limit int64) *MgoSess {
  171. ms.limit = limit
  172. return ms
  173. }
  174. func (ms *MgoSess) Skip(skip int64) *MgoSess {
  175. ms.skip = skip
  176. return ms
  177. }
  178. func (ms *MgoSess) Sort(sorts ...string) *MgoSess {
  179. ms.sorts = sorts
  180. return ms
  181. }
  182. func (ms *MgoSess) Pipe(p []map[string]interface{}) *pipe {
  183. pe := &pipe{
  184. ms: ms,
  185. pipeline: p,
  186. }
  187. return pe
  188. }
  189. func (ms *MgoSess) Insert(doc interface{}) error {
  190. _, err := ms.M.C.Database(ms.db).Collection(ms.coll).InsertOne(ms.M.Ctx, doc)
  191. return err
  192. }
  193. func (ms *MgoSess) Remove(filter interface{}) error {
  194. _, err := ms.M.C.Database(ms.db).Collection(ms.coll).DeleteOne(ms.M.Ctx, filter)
  195. return err
  196. }
  197. func (ms *MgoSess) RemoveId(_id interface{}) error {
  198. _, err := ms.M.C.Database(ms.db).Collection(ms.coll).DeleteOne(ms.M.Ctx, map[string]interface{}{"_id": _id})
  199. return err
  200. }
  201. func (ms *MgoSess) RemoveAll(filter interface{}) (*mongo.DeleteResult, error) {
  202. return ms.M.C.Database(ms.db).Collection(ms.coll).DeleteMany(ms.M.Ctx, filter)
  203. }
  204. func (ms *MgoSess) Upsert(filter, update interface{}) (*mongo.UpdateResult, error) {
  205. ct := options.Update()
  206. ct.SetUpsert(true)
  207. ue := ObjToM(update)
  208. autoUpdateTime(ms.db, ms.coll, ue)
  209. return ms.M.C.Database(ms.db).Collection(ms.coll).UpdateOne(ms.M.Ctx, filter, ue, ct)
  210. }
  211. func (ms *MgoSess) UpsertId(filter, update interface{}) (*mongo.UpdateResult, error) {
  212. ct := options.Update()
  213. ct.SetUpsert(true)
  214. ue := ObjToM(update)
  215. autoUpdateTime(ms.db, ms.coll, ue)
  216. return ms.M.C.Database(ms.db).Collection(ms.coll).UpdateOne(ms.M.Ctx, map[string]interface{}{"_id": filter}, ue, ct)
  217. }
  218. func (ms *MgoSess) UpdateId(filter, update interface{}) error {
  219. ue := ObjToM(update)
  220. autoUpdateTime(ms.db, ms.coll, ue)
  221. _, err := ms.M.C.Database(ms.db).Collection(ms.coll).UpdateOne(ms.M.Ctx, map[string]interface{}{"_id": filter}, ue)
  222. return err
  223. }
  224. func (ms *MgoSess) Update(filter, update interface{}) error {
  225. ue := ObjToM(update)
  226. autoUpdateTime(ms.db, ms.coll, ue)
  227. _, err := ms.M.C.Database(ms.db).Collection(ms.coll).UpdateOne(ms.M.Ctx, filter, ue)
  228. return err
  229. }
  230. func (ms *MgoSess) Count() (int64, error) {
  231. return ms.M.C.Database(ms.db).Collection(ms.coll).CountDocuments(ms.M.Ctx, ms.query)
  232. }
  233. func (ms *MgoSess) One(v *map[string]interface{}) {
  234. of := options.FindOne()
  235. of.SetProjection(ms.fields)
  236. sr := ms.M.C.Database(ms.db).Collection(ms.coll).FindOne(ms.M.Ctx, ms.query, of)
  237. if sr.Err() == nil {
  238. sr.Decode(&v)
  239. }
  240. }
  241. func (ms *MgoSess) All(v *[]map[string]interface{}) {
  242. of := options.Find()
  243. if ms.fields != nil {
  244. of.SetProjection(ms.fields)
  245. }
  246. if len(ms.sorts) > 0 {
  247. of.SetSort(ms.toSort())
  248. }
  249. if ms.skip > 0 {
  250. of.SetSkip(ms.skip)
  251. }
  252. if ms.limit > 0 {
  253. of.SetLimit(ms.limit)
  254. }
  255. cur, err := ms.M.C.Database(ms.db).Collection(ms.coll).Find(ms.M.Ctx, ms.query, of)
  256. if err == nil && cur.Err() == nil {
  257. cur.All(ms.M.Ctx, v)
  258. }
  259. }
  260. func (ms *MgoSess) toSort() bson.D {
  261. sort := bson.D{}
  262. for _, k := range ms.sorts {
  263. switch k[:1] {
  264. case "-":
  265. sort = append(sort, bson.E{k[1:], -1})
  266. case "+":
  267. sort = append(sort, bson.E{k[1:], 1})
  268. default:
  269. sort = append(sort, bson.E{k, 1})
  270. }
  271. }
  272. return sort
  273. }
  274. func (ms *MgoSess) Iter() *MgoIter {
  275. it := &MgoIter{}
  276. coll := ms.M.C.Database(ms.db).Collection(ms.coll)
  277. find := options.Find()
  278. if ms.skip > 0 {
  279. find.SetSkip(ms.skip)
  280. }
  281. if ms.limit > 0 {
  282. find.SetLimit(ms.limit)
  283. }
  284. find.SetBatchSize(100)
  285. if len(ms.sorts) > 0 {
  286. find.SetSort(ms.toSort())
  287. }
  288. if ms.fields != nil {
  289. find.SetProjection(ms.fields)
  290. }
  291. cur, err := coll.Find(ms.M.Ctx, ms.query, find)
  292. if err != nil {
  293. log.Println("mgo find err", err.Error())
  294. }
  295. if err == nil {
  296. it.Cursor = cur
  297. it.Ctx = ms.M.Ctx
  298. }
  299. return it
  300. }
  301. type pipe struct {
  302. ms *MgoSess
  303. pipeline []map[string]interface{}
  304. }
  305. func (p *pipe) All(v *[]map[string]interface{}) {
  306. cur, err := p.ms.M.C.Database(p.ms.db).Collection(p.ms.coll).Aggregate(p.ms.M.Ctx, p.pipeline)
  307. if err == nil && cur.Err() == nil {
  308. cur.All(p.ms.M.Ctx, v)
  309. }
  310. }
  311. func (p *pipe) Iter() *MgoIter {
  312. it := &MgoIter{}
  313. coll := p.ms.M.C.Database(p.ms.db).Collection(p.ms.coll)
  314. aggregate := options.Aggregate()
  315. aggregate.SetBatchSize(100)
  316. cur, err := coll.Aggregate(p.ms.M.Ctx, p.pipeline, aggregate)
  317. if err != nil {
  318. log.Println("mgo aggregate err", err.Error())
  319. }
  320. if err == nil {
  321. it.Cursor = cur
  322. it.Ctx = p.ms.M.Ctx
  323. }
  324. return it
  325. }
  326. type MongodbSim struct {
  327. MongodbAddr string
  328. Size int
  329. // MinSize int
  330. DbName string
  331. C *mongo.Client
  332. Ctx context.Context
  333. ShortCtx context.Context
  334. pool chan bool
  335. UserName string
  336. Password string
  337. ReplSet string
  338. Direct bool
  339. }
  340. func (m *MongodbSim) GetMgoConn() *MgoSess {
  341. //m.Open()
  342. ms := &MgoSess{}
  343. ms.M = m
  344. return ms
  345. }
  346. func (m *MongodbSim) DestoryMongoConn(ms *MgoSess) {
  347. //m.Close()
  348. ms.M = nil
  349. ms = nil
  350. }
  351. func (m *MongodbSim) Destroy() {
  352. //m.Close()
  353. m.C.Disconnect(nil)
  354. m.C = nil
  355. }
  356. func (m *MongodbSim) InitPool() {
  357. opts := options.Client()
  358. if m.Direct {
  359. opts.SetDirect(true)
  360. }
  361. registry := bson.NewRegistryBuilder().RegisterTypeMapEntry(bson.TypeArray, reflect.TypeOf([]interface{}{})).Build()
  362. opts.SetRegistry(registry)
  363. opts.SetConnectTimeout(3 * time.Second)
  364. opts.SetHosts(strings.Split(m.MongodbAddr, ","))
  365. //opts.ApplyURI("mongodb://" + m.MongodbAddr)
  366. opts.SetMaxPoolSize(uint64(m.Size))
  367. if m.UserName != "" && m.Password != "" {
  368. cre := options.Credential{
  369. Username: m.UserName,
  370. Password: m.Password,
  371. }
  372. opts.SetAuth(cre)
  373. }
  374. /*ms := strings.Split(m.MongodbAddr, ",")
  375. if m.ReplSet == "" && len(ms) > 1 {
  376. m.ReplSet = "qfws"
  377. }*/
  378. if m.ReplSet != "" {
  379. opts.SetReplicaSet(m.ReplSet)
  380. opts.SetDirect(false)
  381. }
  382. m.pool = make(chan bool, m.Size)
  383. opts.SetMaxConnIdleTime(2 * time.Hour)
  384. m.Ctx, _ = context.WithTimeout(context.Background(), 99999*time.Hour)
  385. m.ShortCtx, _ = context.WithTimeout(context.Background(), 1*time.Minute)
  386. client, err := mongo.Connect(m.ShortCtx, opts)
  387. if err != nil {
  388. log.Println("mgo init error:", err.Error())
  389. } else {
  390. m.C = client
  391. }
  392. }
  393. func (m *MongodbSim) Open() {
  394. m.pool <- true
  395. }
  396. func (m *MongodbSim) Close() {
  397. <-m.pool
  398. }
  399. //新建表并生成索引
  400. func (m *MongodbSim) CreateIndex(c string, models []mongo.IndexModel) bool {
  401. defer catch()
  402. m.Open()
  403. defer m.Close()
  404. coll := m.C.Database(m.DbName).Collection(c)
  405. names, err := coll.Indexes().CreateMany(m.Ctx, models)
  406. if err == nil && len(names) > 0 {
  407. return true
  408. } else {
  409. log.Println("CreateIndex Error:", err)
  410. return false
  411. }
  412. }
  413. func (m *MongodbSim) Save(c string, doc interface{}) string {
  414. defer catch()
  415. m.Open()
  416. defer m.Close()
  417. coll := m.C.Database(m.DbName).Collection(c)
  418. obj := ObjToM(doc)
  419. id := primitive.NewObjectID()
  420. (*obj)["_id"] = id
  421. _, err := coll.InsertOne(m.Ctx, obj)
  422. if nil != err {
  423. log.Println("SaveError", err)
  424. return ""
  425. }
  426. return id.Hex()
  427. }
  428. //原_id不变
  429. func (m *MongodbSim) SaveByOriID(c string, doc interface{}) bool {
  430. defer catch()
  431. m.Open()
  432. defer m.Close()
  433. coll := m.C.Database(m.DbName).Collection(c)
  434. _, err := coll.InsertOne(m.Ctx, ObjToM(doc))
  435. if nil != err {
  436. log.Println("SaveByOriIDError", err)
  437. return false
  438. }
  439. return true
  440. }
  441. //批量插入
  442. func (m *MongodbSim) SaveBulk(c string, doc ...map[string]interface{}) bool {
  443. defer catch()
  444. m.Open()
  445. defer m.Close()
  446. coll := m.C.Database(m.DbName).Collection(c)
  447. var writes []mongo.WriteModel
  448. for _, d := range doc {
  449. write := mongo.NewInsertOneModel()
  450. write.SetDocument(d)
  451. writes = append(writes, write)
  452. }
  453. br, e := coll.BulkWrite(m.Ctx, writes)
  454. if e != nil {
  455. b := strings.Index(e.Error(), "duplicate") > -1
  456. log.Println("mgo savebulk error:", e.Error())
  457. if br != nil {
  458. log.Println("mgo savebulk size:", br.InsertedCount)
  459. }
  460. return b
  461. }
  462. return true
  463. }
  464. //批量插入
  465. func (m *MongodbSim) SaveBulkInterface(c string, doc ...interface{}) bool {
  466. defer catch()
  467. m.Open()
  468. defer m.Close()
  469. coll := m.C.Database(m.DbName).Collection(c)
  470. var writes []mongo.WriteModel
  471. for _, d := range doc {
  472. write := mongo.NewInsertOneModel()
  473. write.SetDocument(d)
  474. writes = append(writes, write)
  475. }
  476. br, e := coll.BulkWrite(m.Ctx, writes)
  477. if e != nil {
  478. b := strings.Index(e.Error(), "duplicate") > -1
  479. log.Println("mgo SaveBulkInterface error:", e.Error())
  480. if br != nil {
  481. log.Println("mgo SaveBulkInterface size:", br.InsertedCount)
  482. }
  483. return b
  484. }
  485. return true
  486. }
  487. //按条件统计
  488. func (m *MongodbSim) Count(c string, q interface{}) int {
  489. r, _ := m.CountByErr(c, q)
  490. return r
  491. }
  492. //统计
  493. func (m *MongodbSim) CountByErr(c string, q interface{}) (int, error) {
  494. defer catch()
  495. m.Open()
  496. defer m.Close()
  497. var res int64
  498. var err error
  499. if filter := ObjToM(q); filter != nil && len(*filter) > 0 {
  500. res, err = m.C.Database(m.DbName).Collection(c).CountDocuments(m.Ctx, filter)
  501. } else {
  502. res, err = m.C.Database(m.DbName).Collection(c).EstimatedDocumentCount(m.Ctx)
  503. }
  504. if err != nil {
  505. log.Println("统计错误", err.Error())
  506. return 0, err
  507. } else {
  508. return int(res), nil
  509. }
  510. }
  511. //按条件删除
  512. func (m *MongodbSim) Delete(c string, q interface{}) int64 {
  513. defer catch()
  514. m.Open()
  515. defer m.Close()
  516. res, err := m.C.Database(m.DbName).Collection(c).DeleteMany(m.Ctx, ObjToM(q))
  517. if err != nil && res == nil {
  518. log.Println("删除错误", err.Error())
  519. }
  520. return res.DeletedCount
  521. }
  522. //删除对象
  523. func (m *MongodbSim) Del(c string, q interface{}) bool {
  524. defer catch()
  525. m.Open()
  526. defer m.Close()
  527. _, err := m.C.Database(m.DbName).Collection(c).DeleteMany(m.Ctx, ObjToM(q))
  528. if err != nil {
  529. log.Println("删除错误", err.Error())
  530. return false
  531. }
  532. return true
  533. }
  534. //按条件更新
  535. func (m *MongodbSim) Update(c string, q, u interface{}, upsert bool, multi bool) bool {
  536. defer catch()
  537. m.Open()
  538. defer m.Close()
  539. ct := options.Update()
  540. if upsert {
  541. ct.SetUpsert(true)
  542. }
  543. coll := m.C.Database(m.DbName).Collection(c)
  544. ue := ObjToM(u)
  545. autoUpdateTime(m.DbName, c, ue)
  546. var err error
  547. if multi {
  548. _, err = coll.UpdateMany(m.Ctx, ObjToM(q), ue, ct)
  549. } else {
  550. _, err = coll.UpdateOne(m.Ctx, ObjToM(q), ue, ct)
  551. }
  552. if err != nil {
  553. log.Println("更新错误", err.Error())
  554. return false
  555. }
  556. return true
  557. }
  558. func (m *MongodbSim) UpdateById(c string, id interface{}, set interface{}) bool {
  559. defer catch()
  560. m.Open()
  561. defer m.Close()
  562. q := make(map[string]interface{})
  563. if sid, ok := id.(string); ok {
  564. q["_id"], _ = primitive.ObjectIDFromHex(sid)
  565. } else {
  566. q["_id"] = id
  567. }
  568. ue := ObjToM(set)
  569. autoUpdateTime(m.DbName, c, ue)
  570. _, err := m.C.Database(m.DbName).Collection(c).UpdateOne(m.Ctx, q, ue)
  571. if nil != err {
  572. log.Println("UpdateByIdError", err)
  573. return false
  574. }
  575. return true
  576. }
  577. //批量更新
  578. func (m *MongodbSim) UpdateBulkAll(db, c string, doc ...[]map[string]interface{}) bool {
  579. return m.NewUpdateBulk(db, c, false, false, doc...)
  580. }
  581. func (m *MongodbSim) UpdateBulk(c string, doc ...[]map[string]interface{}) bool {
  582. return m.UpdateBulkAll(m.DbName, c, doc...)
  583. }
  584. //批量插入
  585. func (m *MongodbSim) UpSertBulk(c string, doc ...[]map[string]interface{}) bool {
  586. return m.NewUpdateBulk(m.DbName, c, true, false, doc...)
  587. }
  588. //批量插入
  589. func (m *MongodbSim) UpSertMultiBulk(c string, upsert, multi bool, doc ...[]map[string]interface{}) bool {
  590. return m.NewUpdateBulk(m.DbName, c, upsert, multi, doc...)
  591. }
  592. //批量插入
  593. func (m *MongodbSim) NewUpdateBulk(db, c string, upsert, multi bool, doc ...[]map[string]interface{}) bool {
  594. defer catch()
  595. m.Open()
  596. defer m.Close()
  597. coll := m.C.Database(db).Collection(c)
  598. var writes []mongo.WriteModel
  599. for _, d := range doc {
  600. if multi {
  601. write := mongo.NewUpdateManyModel()
  602. write.SetFilter(d[0])
  603. ue := ObjToM(d[1])
  604. autoUpdateTime(m.DbName, c, ue)
  605. write.SetUpdate(ue)
  606. write.SetUpsert(upsert)
  607. writes = append(writes, write)
  608. } else {
  609. write := mongo.NewUpdateOneModel()
  610. write.SetFilter(d[0])
  611. ue := ObjToM(d[1])
  612. autoUpdateTime(m.DbName, c, ue)
  613. write.SetUpdate(ue)
  614. write.SetUpsert(upsert)
  615. writes = append(writes, write)
  616. }
  617. }
  618. br, e := coll.BulkWrite(m.Ctx, writes)
  619. if e != nil {
  620. log.Println("mgo upsert error:", e.Error())
  621. return br == nil || br.UpsertedCount == 0
  622. }
  623. // else {
  624. // if r.UpsertedCount != int64(len(doc)) {
  625. // log.Println("mgo upsert uncomplete:uc/dc", r.UpsertedCount, len(doc))
  626. // }
  627. // return true
  628. // }
  629. return true
  630. }
  631. //查询单条对象
  632. func (m *MongodbSim) FindOne(c string, query interface{}) (*map[string]interface{}, bool) {
  633. return m.FindOneByField(c, query, nil)
  634. }
  635. //查询单条对象
  636. func (m *MongodbSim) FindOneByField(c string, query interface{}, fields interface{}) (*map[string]interface{}, bool) {
  637. defer catch()
  638. res, ok := m.Find(c, query, nil, fields, true, -1, -1)
  639. if nil != res && len(*res) > 0 {
  640. return &((*res)[0]), ok
  641. }
  642. return nil, ok
  643. }
  644. //查询单条对象
  645. func (m *MongodbSim) FindById(c string, query string, fields interface{}) (*map[string]interface{}, bool) {
  646. defer catch()
  647. m.Open()
  648. defer m.Close()
  649. of := options.FindOne()
  650. of.SetProjection(ObjToOth(fields))
  651. res := make(map[string]interface{})
  652. _id, err := primitive.ObjectIDFromHex(query)
  653. if err != nil {
  654. log.Println("_id error", err)
  655. return &res, true
  656. }
  657. sr := m.C.Database(m.DbName).Collection(c).FindOne(m.Ctx, map[string]interface{}{"_id": _id}, of)
  658. if sr.Err() == nil {
  659. sr.Decode(&res)
  660. }
  661. return &res, true
  662. }
  663. //底层查询方法
  664. func (m *MongodbSim) Find(c string, query interface{}, order interface{}, fields interface{}, single bool, start int, limit int) (*[]map[string]interface{}, bool) {
  665. defer catch()
  666. m.Open()
  667. defer m.Close()
  668. var res []map[string]interface{}
  669. coll := m.C.Database(m.DbName).Collection(c)
  670. if single {
  671. res = make([]map[string]interface{}, 1)
  672. of := options.FindOne()
  673. of.SetProjection(ObjToOth(fields))
  674. of.SetSort(ObjToM(order))
  675. if sr := coll.FindOne(m.Ctx, ObjToM(query), of); sr.Err() == nil {
  676. sr.Decode(&res[0])
  677. }
  678. } else {
  679. res = []map[string]interface{}{}
  680. of := options.Find()
  681. of.SetProjection(ObjToOth(fields))
  682. of.SetSort(ObjToM(order))
  683. if start > -1 {
  684. of.SetSkip(int64(start))
  685. of.SetLimit(int64(limit))
  686. }
  687. cur, err := coll.Find(m.Ctx, ObjToM(query), of)
  688. if err == nil && cur.Err() == nil {
  689. cur.All(m.Ctx, &res)
  690. }
  691. }
  692. return &res, true
  693. }
  694. func ObjToOth(query interface{}) *bson.M {
  695. return ObjToMQ(query, false)
  696. }
  697. func ObjToM(query interface{}) *bson.M {
  698. return ObjToMQ(query, true)
  699. }
  700. //obj(string,M)转M,查询用到
  701. func ObjToMQ(query interface{}, isQuery bool) *bson.M {
  702. data := make(bson.M)
  703. defer catch()
  704. if s2, ok2 := query.(*map[string]interface{}); ok2 {
  705. data = bson.M(*s2)
  706. } else if s3, ok3 := query.(*bson.M); ok3 {
  707. return s3
  708. } else if s3, ok3 := query.(*primitive.M); ok3 {
  709. return s3
  710. } else if s, ok := query.(string); ok {
  711. json.Unmarshal([]byte(strings.Replace(s, "'", "\"", -1)), &data)
  712. if ss, oks := data["_id"]; oks && isQuery {
  713. switch ss.(type) {
  714. case string:
  715. data["_id"], _ = primitive.ObjectIDFromHex(ss.(string))
  716. case map[string]interface{}:
  717. tmp := ss.(map[string]interface{})
  718. for k, v := range tmp {
  719. tmp[k], _ = primitive.ObjectIDFromHex(v.(string))
  720. }
  721. data["_id"] = tmp
  722. }
  723. }
  724. } else if s1, ok1 := query.(map[string]interface{}); ok1 {
  725. data = s1
  726. } else if s4, ok4 := query.(bson.M); ok4 {
  727. data = s4
  728. } else if s4, ok4 := query.(primitive.M); ok4 {
  729. data = s4
  730. } else {
  731. data = nil
  732. }
  733. return &data
  734. }
  735. func intAllDef(num interface{}, defaultNum int) int {
  736. if i, ok := num.(int); ok {
  737. return int(i)
  738. } else if i0, ok0 := num.(int32); ok0 {
  739. return int(i0)
  740. } else if i1, ok1 := num.(float64); ok1 {
  741. return int(i1)
  742. } else if i2, ok2 := num.(int64); ok2 {
  743. return int(i2)
  744. } else if i3, ok3 := num.(float32); ok3 {
  745. return int(i3)
  746. } else if i4, ok4 := num.(string); ok4 {
  747. in, _ := strconv.Atoi(i4)
  748. return int(in)
  749. } else if i5, ok5 := num.(int16); ok5 {
  750. return int(i5)
  751. } else if i6, ok6 := num.(int8); ok6 {
  752. return int(i6)
  753. } else if i7, ok7 := num.(*big.Int); ok7 {
  754. in, _ := strconv.Atoi(fmt.Sprint(i7))
  755. return int(in)
  756. } else if i8, ok8 := num.(*big.Float); ok8 {
  757. in, _ := strconv.Atoi(fmt.Sprint(i8))
  758. return int(in)
  759. } else {
  760. return defaultNum
  761. }
  762. }
  763. //出错拦截
  764. func catch() {
  765. if r := recover(); r != nil {
  766. log.Println(r)
  767. for skip := 0; ; skip++ {
  768. _, file, line, ok := runtime.Caller(skip)
  769. if !ok {
  770. break
  771. }
  772. go log.Printf("%v,%v\n", file, line)
  773. }
  774. }
  775. }
  776. //根据bsonID转string
  777. func BsonIdToSId(uid interface{}) string {
  778. if uid == nil {
  779. return ""
  780. } else if u, ok := uid.(string); ok {
  781. return u
  782. } else if u, ok := uid.(primitive.ObjectID); ok {
  783. return u.Hex()
  784. } else {
  785. return ""
  786. }
  787. }
  788. func StringTOBsonId(id string) (bid primitive.ObjectID) {
  789. defer catch()
  790. if id != "" {
  791. bid, _ = primitive.ObjectIDFromHex(id)
  792. }
  793. return
  794. }
  795. func ToObjectIds(ids []string) []primitive.ObjectID {
  796. _ids := []primitive.ObjectID{}
  797. for _, v := range ids {
  798. _id, _ := primitive.ObjectIDFromHex(v)
  799. _ids = append(_ids, _id)
  800. }
  801. return _ids
  802. }
  803. //自动添加更新时间
  804. func autoUpdateTime(db, coll string, ue *bson.M) {
  805. if coll == "user" {
  806. set := ObjToM((*ue)["$set"])
  807. if *set == nil {
  808. set = &bson.M{}
  809. }
  810. (*set)["auto_updatetime"] = time.Now().Unix()
  811. (*ue)["$set"] = set
  812. }
  813. }
  814. func IsObjectIdHex(hex string) bool {
  815. _, err := primitive.ObjectIDFromHex(hex)
  816. if err != nil {
  817. return false
  818. }
  819. return true
  820. }