mongodb.go 21 KB

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