更新 MongoDB 插入更新函数

elastic/elasticSim.go

@@ -525,3 +525,24 @@ func (e *Elastic) DeleteByID(index, id string) error {
 
 	return err
 }
+
+// UpdateDocument 更新指定ID的文档
+func (e *Elastic) UpdateDocument(indexName string, documentID string, updateData map[string]interface{}) error {
+	client := e.GetEsConn()
+	defer e.DestoryEsConn(client)
+	updateResult, err := client.Update().
+		Index(indexName).
+		Id(documentID).
+		Doc(updateData).
+		Do(context.Background())
+
+	if err != nil {
+		return err
+	}
+
+	if updateResult.Result != "updated" {
+		return fmt.Errorf("Document not updated: %v", updateResult.Result)
+	}
+
+	return nil
+}

elastic/elastic_test.go

@@ -22,10 +22,10 @@ var esClinet Elastic
 //}
 func TestCount(t *testing.T) {
 	esClinet = Elastic{
-		S_esurl:  "http://127.0.0.1:19805",
+		S_esurl:  "http://192.168.3.241:40801",
 		I_size:   2,
-		Username: "es_all",
-		Password: "TopJkO2E_d1x",
+		Username: "",
+		Password: "",
 	}
 	esClinet.InitElasticSize()
 
@@ -81,7 +81,7 @@ func TestElastic_ExistsIndex(t *testing.T) {
 	}
 	esClinet.InitElasticSize()
 
-	res, err := esClinet.ExistsIndex("bidding_20231118")
+	res, err := esClinet.ExistsIndex("bidding")
 	assert.Equal(t, nil, err)
 	assert.Equal(t, true, res)
 }
@@ -140,6 +140,6 @@ func TestElastic_DeleteByID(t *testing.T) {
 	}
 	esClinet.InitElasticSize()
 
-	err := esClinet.DeleteByID("bidding_202311", "655c99c60687916fae58eec1")
+	err := esClinet.DeleteByID("bidding_2023112216", "655c99c60687916fae58eec1")
 	assert.Equal(t, nil, err)
 }

mongodb/mgo_test.go

@@ -3,6 +3,7 @@ package mongodb
 import (
 	"fmt"
 	"testing"
+	"time"
 )
 
 var mgo *MongodbSim
@@ -38,3 +39,13 @@ func TestMongodbSim_Save(t *testing.T) {
 	res := mgo.Save("wcc_wcc", data)
 	fmt.Println(res)
 }
+
+func TestMongodbSim_InsertOrUpdate(t *testing.T) {
+	data := map[string]interface{}{
+		"_id":        StringTOBsonId("5f289ea352c1d9fbf84e1f1b"),
+		"num":        12,
+		"createtime": time.Now().Unix(),
+	}
+	err := mgo.InsertOrUpdate("wcc", "bidding_hot", data)
+	fmt.Println(err)
+}

mongodb/mongodb.go

@@ -562,6 +562,8 @@ func (m *MongodbSim) Update(c string, q, u interface{}, upsert bool, multi bool)
 	}
 	return true
 }
+
+//InsertOrUpdate 插入或更新
 func (m *MongodbSim) InsertOrUpdate(db, collectionName string, data map[string]interface{}) error {
 	defer catch()
 	m.Open()
@@ -570,31 +572,39 @@ func (m *MongodbSim) InsertOrUpdate(db, collectionName string, data map[string]i
 	idValue, ok := data["_id"]
 	var idObj primitive.ObjectID
 	if ok {
-		if sid, ok := idValue.(string); ok {
-			idValue, _ = primitive.ObjectIDFromHex(sid)
+		if oid, ok := idValue.(primitive.ObjectID); ok {
+			idObj = oid
+		} else if sid, ok := idValue.(string); ok {
+			idObj, _ = primitive.ObjectIDFromHex(sid)
 		}
-	} else {
-		// 如果 _id 不存在，则生成一个新的 ObjectID
+	}
+
+	// 如果 _id 不存在，则生成一个新的 ObjectID
+	if idObj.IsZero() {
 		idObj = primitive.NewObjectID()
-		data["_id"] = idObj.Hex()
+		data["_id"] = idObj
 	}
 
-	// 根据 _id 字段进行查找和更新，如果不存在则插入新数据
+	// 在插入新文档时检查 _id 是否已存在
 	filter := bson.M{"_id": idObj}
-	update := bson.M{"$set": data}
-
-	// 尝试更新记录，如果记录不存在则插入新数据
-	updateResult, err := m.C.Database(db).Collection(collectionName).UpdateOne(context.Background(), filter, update, options.Update().SetUpsert(true))
+	count, err := m.C.Database(db).Collection(collectionName).CountDocuments(context.Background(), filter)
 	if err != nil {
 		return err
 	}
 
-	// 如果未修改任何文档，则表示需要插入新数据
-	if updateResult.ModifiedCount == 0 {
+	// 如果 _id 不存在，则进行插入
+	if count == 0 {
 		_, err := m.C.Database(db).Collection(collectionName).InsertOne(context.Background(), data)
 		if err != nil {
 			return err
 		}
+	} else {
+		// 如果 _id 已存在，则进行更新
+		update := bson.M{"$set": data}
+		_, err := m.C.Database(db).Collection(collectionName).UpdateOne(context.Background(), filter, update)
+		if err != nil {
+			return err
+		}
 	}
 
 	return nil