Open-IM-Server/pkg/common/db/unrelation/super_group.go

package unrelation

import (
	"Open_IM/pkg/common/config"
	"Open_IM/pkg/common/db/table/unrelation"
	"Open_IM/pkg/utils"
	"context"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readconcern"
)

var _ unrelation.SuperGroupModelInterface = (*SuperGroupMongoDriver)(nil)

func NewSuperGroupMongoDriver(mgoClient *mongo.Client) *SuperGroupMongoDriver {
	mgoDB := mgoClient.Database(config.Config.Mongo.DBDatabase)
	return &SuperGroupMongoDriver{MgoDB: mgoDB, MgoClient: mgoClient, superGroupCollection: mgoDB.Collection(unrelation.CSuperGroup), userToSuperGroupCollection: mgoDB.Collection(unrelation.CUserToSuperGroup)}
}

type SuperGroupMongoDriver struct {
	MgoClient                  *mongo.Client
	MgoDB                      *mongo.Database
	superGroupCollection       *mongo.Collection
	userToSuperGroupCollection *mongo.Collection
}

//	func (s *SuperGroupMongoDriver) CreateSuperGroup(ctx context.Context, groupID string, initMemberIDs []string, tx ...interface{}) error {
//			superGroup := unrelation.SuperGroupModel{
//				GroupID:   groupID,
//				MemberIDs: initMemberIDs,
//			}
//			coll := getTxCtx(s.superGroupCollection, tx)
//			_, err := coll.InsertOne(ctx, superGroup)
//			if err != nil {
//				return err
//			}
//			opts := &options.UpdateOptions{
//				Upsert: utils.ToPtr(true),
//			}
//			for _, userID := range initMemberIDs {
//				_, err = coll.UpdateOne(ctx, bson.M{"user_id": userID}, bson.M{"$addToSet": bson.M{"group_id_list": groupID}}, opts)
//				if err != nil {
//					return err
//				}
//			}
//			return nil
//	}
//
//	func (s *SuperGroupMongoDriver) FindSuperGroup(ctx context.Context, groupIDs []string, tx ...interface{}) (groups []*unrelation.SuperGroupModel, err error) {
//		cursor, err := s.superGroupCollection.Find(ctx, bson.M{"group_id": bson.M{
//			"$in": groupIDs,
//		}})
//		if err != nil {
//			return nil, utils.Wrap(err, "")
//		}
//		defer cursor.Close(ctx)
//		if err := cursor.All(ctx, &groups); err != nil {
//			return nil, utils.Wrap(err, "")
//		}
//		return groups, nil
//	}
//
//	func (s *SuperGroupMongoDriver) AddUserToSuperGroup(ctx context.Context, groupID string, userIDs []string, tx ...interface{}) error {
//			opts := options.Session().SetDefaultReadConcern(readconcern.Majority())
//			return s.MgoDB.Client().UseSessionWithOptions(ctx, opts, func(sCtx mongo.SessionContext) error {
//				_, err := s.superGroupCollection.UpdateOne(sCtx, bson.M{"group_id": groupID}, bson.M{"$addToSet": bson.M{"member_id_list": bson.M{"$each": userIDs}}})
//				if err != nil {
//					_ = sCtx.AbortTransaction(ctx)
//					return err
//				}
//				upsert := true
//				opts := &options.UpdateOptions{
//					Upsert: &upsert,
//				}
//				for _, userID := range userIDs {
//					_, err = s.userToSuperGroupCollection.UpdateOne(sCtx, bson.M{"user_id": userID}, bson.M{"$addToSet": bson.M{"group_id_list": groupID}}, opts)
//					if err != nil {
//						_ = sCtx.AbortTransaction(ctx)
//						return utils.Wrap(err, "transaction failed")
//					}
//				}
//				return sCtx.CommitTransaction(ctx)
//			})
//	}
//
//	func (s *SuperGroupMongoDriver) RemoverUserFromSuperGroup(ctx context.Context, groupID string, userIDs []string, tx ...interface{}) error {
//			opts := options.Session().SetDefaultReadConcern(readconcern.Majority())
//			return s.MgoDB.Client().UseSessionWithOptions(ctx, opts, func(sCtx mongo.SessionContext) error {
//				_, err := s.superGroupCollection.UpdateOne(sCtx, bson.M{"group_id": groupID}, bson.M{"$pull": bson.M{"member_id_list": bson.M{"$in": userIDs}}})
//				if err != nil {
//					_ = sCtx.AbortTransaction(ctx)
//					return err
//				}
//				err = s.RemoveGroupFromUser(sCtx, groupID, userIDs)
//				if err != nil {
//					_ = sCtx.AbortTransaction(ctx)
//					return err
//				}
//				return sCtx.CommitTransaction(ctx)
//			})
//	}
//
//	func (s *SuperGroupMongoDriver) GetSuperGroupByUserID(ctx context.Context, userID string, tx ...interface{}) (*unrelation.UserToSuperGroupModel, error) {
//		//TODO implement me
//		panic("implement me")
//	}
//
//	func (s *SuperGroupMongoDriver) DeleteSuperGroup(ctx context.Context, groupID string, tx ...interface{}) error {
//		//TODO implement me
//		panic("implement me")
//	}

func (s *SuperGroupMongoDriver) Transaction(ctx context.Context, fn func(s unrelation.SuperGroupModelInterface, tx any) error) error {
	sess, err := s.MgoClient.StartSession()
	if err != nil {
		return err
	}
	txCtx := mongo.NewSessionContext(ctx, sess)
	defer sess.EndSession(txCtx)
	if err := fn(s, txCtx); err != nil {
		_ = sess.AbortTransaction(txCtx)
		return err
	}
	return utils.Wrap(sess.CommitTransaction(txCtx), "")
}

func (s *SuperGroupMongoDriver) getTxCtx(ctx context.Context, tx []any) context.Context {
	if len(tx) > 0 {
		if ctx, ok := tx[0].(mongo.SessionContext); ok {
			return ctx
		}
	}
	return ctx
}

//func (s *SuperGroupMongoDriver) Transaction(ctx context.Context, fn func(ctx mongo.SessionContext) error) error {
//	sess, err := s.MgoClient.StartSession()
//	if err != nil {
//		return err
//	}
//	sCtx := mongo.NewSessionContext(ctx, sess)
//
//	defer sess.EndSession(sCtx)
//	if err := fn(sCtx); err != nil {
//		_ = sess.AbortTransaction(sCtx)
//		return err
//	}
//	return utils.Wrap(sess.CommitTransaction(sCtx), "")
//}

func (s *SuperGroupMongoDriver) CreateSuperGroup(ctx context.Context, groupID string, initMemberIDs []string, tx ...any) error {
	ctx = s.getTxCtx(ctx, tx)
	_, err := s.superGroupCollection.InsertOne(ctx, &unrelation.SuperGroupModel{
		GroupID:   groupID,
		MemberIDs: initMemberIDs,
	})
	if err != nil {
		return err
	}
	for _, userID := range initMemberIDs {
		_, err = s.userToSuperGroupCollection.UpdateOne(ctx, bson.M{"user_id": userID}, bson.M{"$addToSet": bson.M{"group_id_list": groupID}}, &options.UpdateOptions{
			Upsert: utils.ToPtr(true),
		})
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *SuperGroupMongoDriver) TakeSuperGroup(ctx context.Context, groupID string, tx ...any) (group *unrelation.SuperGroupModel, err error) {
	ctx = s.getTxCtx(ctx, tx)
	if err := s.superGroupCollection.FindOne(ctx, bson.M{"group_id": groupID}).Decode(&group); err != nil {
		return nil, utils.Wrap(err, "")
	}
	return group, nil
}

func (s *SuperGroupMongoDriver) FindSuperGroup(ctx context.Context, groupIDs []string, tx ...any) (groups []*unrelation.SuperGroupModel, err error) {
	ctx = s.getTxCtx(ctx, tx)
	cursor, err := s.superGroupCollection.Find(ctx, bson.M{"group_id": bson.M{
		"$in": groupIDs,
	}})
	if err != nil {
		return nil, utils.Wrap(err, "")
	}
	defer cursor.Close(ctx)
	if err := cursor.All(ctx, &groups); err != nil {
		return nil, utils.Wrap(err, "")
	}
	return groups, nil
}

func (s *SuperGroupMongoDriver) AddUserToSuperGroup(ctx context.Context, groupID string, userIDs []string, tx ...any) error {
	ctx = s.getTxCtx(ctx, tx)
	opts := options.Session().SetDefaultReadConcern(readconcern.Majority())
	return s.MgoDB.Client().UseSessionWithOptions(ctx, opts, func(sCtx mongo.SessionContext) error {
		_, err := s.superGroupCollection.UpdateOne(sCtx, bson.M{"group_id": groupID}, bson.M{"$addToSet": bson.M{"member_id_list": bson.M{"$each": userIDs}}})
		if err != nil {
			_ = sCtx.AbortTransaction(ctx)
			return err
		}
		upsert := true
		opts := &options.UpdateOptions{
			Upsert: &upsert,
		}
		for _, userID := range userIDs {
			_, err = s.userToSuperGroupCollection.UpdateOne(sCtx, bson.M{"user_id": userID}, bson.M{"$addToSet": bson.M{"group_id_list": groupID}}, opts)
			if err != nil {
				_ = sCtx.AbortTransaction(ctx)
				return utils.Wrap(err, "transaction failed")
			}
		}
		return sCtx.CommitTransaction(ctx)
	})
}

func (s *SuperGroupMongoDriver) RemoverUserFromSuperGroup(ctx context.Context, groupID string, userIDs []string, tx ...any) error {
	ctx = s.getTxCtx(ctx, tx)
	opts := options.Session().SetDefaultReadConcern(readconcern.Majority())
	return s.MgoDB.Client().UseSessionWithOptions(ctx, opts, func(sCtx mongo.SessionContext) error {
		_, err := s.superGroupCollection.UpdateOne(sCtx, bson.M{"group_id": groupID}, bson.M{"$pull": bson.M{"member_id_list": bson.M{"$in": userIDs}}})
		if err != nil {
			_ = sCtx.AbortTransaction(ctx)
			return err
		}
		err = s.RemoveGroupFromUser(sCtx, groupID, userIDs)
		if err != nil {
			_ = sCtx.AbortTransaction(ctx)
			return err
		}
		return sCtx.CommitTransaction(ctx)
	})
}

func (s *SuperGroupMongoDriver) GetSuperGroupByUserID(ctx context.Context, userID string, tx ...any) (*unrelation.UserToSuperGroupModel, error) {
	ctx = s.getTxCtx(ctx, tx)
	var user unrelation.UserToSuperGroupModel
	err := s.userToSuperGroupCollection.FindOne(ctx, bson.M{"user_id": userID}).Decode(&user)
	return &user, utils.Wrap(err, "")
}

func (s *SuperGroupMongoDriver) DeleteSuperGroup(ctx context.Context, groupID string, tx ...any) error {
	ctx = s.getTxCtx(ctx, tx)
	group, err := s.TakeSuperGroup(ctx, groupID, tx...)
	if err != nil {
		return err
	}
	if _, err := s.superGroupCollection.DeleteOne(ctx, bson.M{"group_id": groupID}); err != nil {
		return utils.Wrap(err, "")
	}
	return s.RemoveGroupFromUser(ctx, groupID, group.MemberIDs)
}

//func (s *SuperGroupMongoDriver) DeleteSuperGroup(ctx context.Context, groupID string, tx ...any) error {
//	ctx = s.getTxCtx(ctx, tx)
//	opts := options.Session().SetDefaultReadConcern(readconcern.Majority())
//	return s.MgoDB.Client().UseSessionWithOptions(ctx, opts, func(sCtx mongo.SessionContext) error {
//		superGroup := &unrelation.SuperGroupModel{}
//		_, err := s.superGroupCollection.DeleteOne(sCtx, bson.M{"group_id": groupID})
//		if err != nil {
//			_ = sCtx.AbortTransaction(ctx)
//			return err
//		}
//		if err = s.RemoveGroupFromUser(sCtx, groupID, superGroup.MemberIDs); err != nil {
//			_ = sCtx.AbortTransaction(ctx)
//			return err
//		}
//		return sCtx.CommitTransaction(ctx)
//	})
//}

func (s *SuperGroupMongoDriver) RemoveGroupFromUser(ctx context.Context, groupID string, userIDs []string, tx ...any) error {
	ctx = s.getTxCtx(ctx, tx)
	_, err := s.userToSuperGroupCollection.UpdateOne(ctx, bson.M{"user_id": bson.M{"$in": userIDs}}, bson.M{"$pull": bson.M{"group_id_list": groupID}})
	return utils.Wrap(err, "")
}