复制项目

pkg/common/storage/database/mgo/seq_conversation.go

@@ -0,0 +1,104 @@
+package mgo
+
+import (
+	"context"
+	"errors"
+
+	"git.imall.cloud/openim/open-im-server-deploy/pkg/common/storage/database"
+	"git.imall.cloud/openim/open-im-server-deploy/pkg/common/storage/model"
+	"github.com/openimsdk/tools/db/mongoutil"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+func NewSeqConversationMongo(db *mongo.Database) (database.SeqConversation, error) {
+	coll := db.Collection(database.SeqConversationName)
+	_, err := coll.Indexes().CreateOne(context.Background(), mongo.IndexModel{
+		Keys: bson.D{
+			{Key: "conversation_id", Value: 1},
+		},
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &seqConversationMongo{coll: coll}, nil
+}
+
+type seqConversationMongo struct {
+	coll *mongo.Collection
+}
+
+func (s *seqConversationMongo) setSeq(ctx context.Context, conversationID string, seq int64, field string) error {
+	filter := map[string]any{
+		"conversation_id": conversationID,
+	}
+	insert := bson.M{
+		"conversation_id": conversationID,
+		"min_seq":         0,
+		"max_seq":         0,
+	}
+	delete(insert, field)
+	update := map[string]any{
+		"$set": bson.M{
+			field: seq,
+		},
+		"$setOnInsert": insert,
+	}
+	opt := options.Update().SetUpsert(true)
+	return mongoutil.UpdateOne(ctx, s.coll, filter, update, false, opt)
+}
+
+func (s *seqConversationMongo) Malloc(ctx context.Context, conversationID string, size int64) (int64, error) {
+	if size < 0 {
+		return 0, errors.New("size must be greater than 0")
+	}
+	if size == 0 {
+		return s.GetMaxSeq(ctx, conversationID)
+	}
+	filter := map[string]any{"conversation_id": conversationID}
+	update := map[string]any{
+		"$inc": map[string]any{"max_seq": size},
+		"$set": map[string]any{"min_seq": int64(0)},
+	}
+	opt := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After).SetProjection(map[string]any{"_id": 0, "max_seq": 1})
+	lastSeq, err := mongoutil.FindOneAndUpdate[int64](ctx, s.coll, filter, update, opt)
+	if err != nil {
+		return 0, err
+	}
+	return lastSeq - size, nil
+}
+
+func (s *seqConversationMongo) SetMaxSeq(ctx context.Context, conversationID string, seq int64) error {
+	return s.setSeq(ctx, conversationID, seq, "max_seq")
+}
+
+func (s *seqConversationMongo) GetMaxSeq(ctx context.Context, conversationID string) (int64, error) {
+	seq, err := mongoutil.FindOne[int64](ctx, s.coll, bson.M{"conversation_id": conversationID}, options.FindOne().SetProjection(map[string]any{"_id": 0, "max_seq": 1}))
+	if err == nil {
+		return seq, nil
+	} else if IsNotFound(err) {
+		return 0, nil
+	} else {
+		return 0, err
+	}
+}
+
+func (s *seqConversationMongo) GetMinSeq(ctx context.Context, conversationID string) (int64, error) {
+	seq, err := mongoutil.FindOne[int64](ctx, s.coll, bson.M{"conversation_id": conversationID}, options.FindOne().SetProjection(map[string]any{"_id": 0, "min_seq": 1}))
+	if err == nil {
+		return seq, nil
+	} else if IsNotFound(err) {
+		return 0, nil
+	} else {
+		return 0, err
+	}
+}
+
+func (s *seqConversationMongo) SetMinSeq(ctx context.Context, conversationID string, seq int64) error {
+	return s.setSeq(ctx, conversationID, seq, "min_seq")
+}
+
+func (s *seqConversationMongo) GetConversation(ctx context.Context, conversationID string) (*model.SeqConversation, error) {
+	return mongoutil.FindOne[*model.SeqConversation](ctx, s.coll, bson.M{"conversation_id": conversationID})
+}