Go 封装 mongo

一、封装代码

package mongodb

import (
	"context"
	"log"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

var client *mongo.Client

func NewMongoDB(ctx context.Context, appUri string) error {
	//连接到mongodb
	c, err := mongo.Connect(ctx, options.Client().ApplyURI(appUri))
	if err != nil {
		return err
	}
	//检查连接
	err = c.Ping(ctx, nil)
	if err != nil {
		return err
	}
	client = c
	log.Println("mongodb连接成功")
	return nil
}

type MongoDB[T any] struct {
	database   string
	collection string
}

func NewMGDB[T any](database string, collection string) *MongoDB[T] {
	if client == nil {
		log.Fatalln("mongo.Client Not initialized !")
	}
	return &MongoDB[T]{
		database,
		collection,
	}
}

//新增一条记录
func (mg *MongoDB[T]) InsertOne(ctx context.Context, value T) *mongo.InsertOneResult {
	result, err := mg.getCollection().InsertOne(ctx, value)
	if err != nil {
		panic(err)
	}
	return result
}

//新增多条记录
func (mg *MongoDB[T]) InsertMultiple(ctx context.Context, data []T) *mongo.InsertManyResult {
	var array []interface{}
	for i := 0; i < len(data); i++ {
		array = append(array, data[i])
	}
	result, err := mg.getCollection().InsertMany(ctx, array)
	if err != nil {
		panic(err)
	}
	return result
}

//根据字段名和值查询一条记录
func (mg *MongoDB[T]) FindOne(ctx context.Context, filter filter) (T, error) {
	var t T
	err := mg.getCollection().FindOne(ctx, filter).Decode(&t)
	if err != nil {
		return t, err
	}
	return t, nil
}

//根据条件查询多条记录
func (mg *MongoDB[T]) Find(ctx context.Context, filter filter, limit int64) ([]T, error) {
	findOpts := options.Find()
	findOpts.SetLimit(limit)
	cursor, err := mg.getCollection().Find(ctx, filter, findOpts)
	var ts []T
	if err != nil {
		return ts, err
	}
	for cursor.Next(ctx) {
		var t T
		err := cursor.Decode(&t)
		if err != nil {
			return ts, err
		}
		ts = append(ts, t)
	}
	cursor.Close(ctx)
	return ts, nil
}

//根据条件更新
func (mg *MongoDB[T]) UpdateOne(ctx context.Context, filter filter, update interface{}) (int64, error) {
	result, err := mg.getCollection().UpdateOne(ctx, filter, bson.M{"$set": update})
	return result.ModifiedCount, err
}

//根据id更新
func (mg *MongoDB[T]) UpdateOneById(ctx context.Context, id string, update interface{}) (int64, error) {
	result, err := mg.getCollection().UpdateOne(ctx, filter{{Key: "_id", Value: mg.ObjectID(id)}}, update)
	return result.ModifiedCount, err
}

//更新多个
func (mg *MongoDB[T]) UpdateMany(ctx context.Context, filter filter, update interface{}) (int64, error) {
	result, err := mg.getCollection().UpdateMany(ctx, filter, bson.D{{Key: "$set", Value: update}})
	return result.ModifiedCount, err
}

//获取表
func (mg *MongoDB[T]) getCollection() *mongo.Collection {
	return client.Database(mg.database).Collection(mg.collection)
}

//删除一条记录
func (mg *MongoDB[T]) DeleteOne(ctx context.Context, filter filter) (int64, error) {
	result, err := mg.getCollection().DeleteOne(ctx, filter)
	return result.DeletedCount, err
}

//根据id删除一条记录
func (mg *MongoDB[T]) DeleteOneById(ctx context.Context, id string) (int64, error) {
	result, err := mg.getCollection().DeleteOne(ctx, filter{{Key: "_id", Value: mg.ObjectID(id)}})
	return result.DeletedCount, err
}

//删除多条记录
func (mg *MongoDB[T]) DeleteMany(ctx context.Context, filter filter) (int64, error) {
	result, err := mg.getCollection().DeleteMany(ctx, filter)
	return result.DeletedCount, err
}

//objcetid
func (mg *MongoDB[T]) ObjectID(id string) primitive.ObjectID {
	objectId, err := primitive.ObjectIDFromHex(id)
	if err != nil {
		log.Fatal(err)
	}
	return objectId
}

//定义过滤器
type filter bson.D

//匹配字段值大于指定值的文档
func (f filter) GT(key string, value interface{}) filter {
	f = append(f, bson.E{Key: key, Value: bson.D{{Key: "$gt", Value: value}}})
	return f
}

//匹配字段值大于等于指定值的文档
func (f filter) GTE(key string, value interface{}) filter {
	f = append(f, bson.E{Key: key, Value: bson.D{{Key: "$gte", Value: value}}})
	return f
}

//匹配字段值等于指定值的文档
func (f filter) EQ(key string, value interface{}) filter {
	f = append(f, bson.E{Key: key, Value: bson.D{{Key: "$eq", Value: value}}})
	return f
}

//匹配字段值小于指定值的文档
func (f filter) LT(key string, value interface{}) filter {
	f = append(f, bson.E{Key: key, Value: bson.D{{Key: "$lt", Value: value}}})
	return f
}

//匹配字段值小于等于指定值的文档
func (f filter) LET(key string, value interface{}) filter {
	f = append(f, bson.E{Key: key, Value: bson.D{{Key: "$let", Value: value}}})
	return f
}

//匹配字段值不等于指定值的文档，包括没有这个字段的文档
func (f filter) NE(key string, value interface{}) filter {
	f = append(f, bson.E{Key: key, Value: bson.D{{Key: "$ne", Value: value}}})
	return f
}

//匹配字段值等于指定数组中的任何值
func (f filter) IN(key string, value ...interface{}) filter {
	f = append(f, bson.E{Key: key, Value: bson.D{{Key: "$in", Value: value}}})
	return f
}

//字段值不在指定数组或者不存在
func (f filter) NIN(key string, value ...interface{}) filter {
	f = append(f, bson.E{Key: key, Value: bson.D{{Key: "$nin", Value: value}}})
	return f
}

//创建一个条件查询对象
func Newfilter() filter {
	return filter{}
}

二、测试代码

type Test struct {
	Id            primitive.ObjectID `bson:"_id"`
	Title         string             `bson:"title"`
	Author        string             `bson:"author"`
	YearPublished int64              `bson:"year_published"`
}

func main() {
	ctx := context.Background()

	//连接数据库
	err := mongodb.NewMongoDB(ctx, "mongodb://localhost:27017")
	if err != nil {
		log.Fatalf("%s", err)
		return

	}

	//设置使用的库和表
	mgdb := mongodb.NewMGDB[Test]("test", "favorite_books")

	//插入单条
	insertOneResult := mgdb.InsertOne(ctx, Test{
		Id:            primitive.NewObjectID(),
		Title:         "test",
		Author:        "author test",
		YearPublished: 9999,
	})

	log.Printf("插入单条记录: %v \n", insertOneResult.InsertedID)

	//插入多条
	var tests []Test
	for i := 1; i < 100; i++ {
		tests = append(tests, Test{
			Id:            primitive.NewObjectID(),
			Title:         "test_" + fmt.Sprintf("%d", i),
			Author:        "author test " + fmt.Sprintf("%d", i),
			YearPublished: int64(i),
		})
	}
	insertMultipleResult := mgdb.InsertMultiple(ctx, tests)

	log.Printf("插入多条记录: %v \n", insertMultipleResult.InsertedIDs)

	//查询
	filter := mongodb.Newfilter().EQ("title", "test").EQ("author", "author test")
	result, err := mgdb.FindOne(ctx, filter)
	if err != nil {
		log.Fatalf("%s", err)
	}
	buf, err := json.Marshal(result)
	fmt.Printf("查询单条记录: %s\n  ", string(buf))

	//查询
	filter = mongodb.Newfilter().GT("year_published", 5).LT("year_published", 10)
	results, err := mgdb.Find(ctx, filter, 10)
	if err != nil {
		log.Fatalf("%s", err)
	}
	buf, err = json.Marshal(results)
	fmt.Printf("查询多条记录: %v\n  ", string(buf))

	//单条记录更新
	filter = mongodb.Newfilter().EQ("year_published", 9999)
	updateCount, err := mgdb.UpdateOne(ctx, filter, map[string]interface{}{
		"author": "test 00021",
	})
	if err != nil {
		log.Fatalf("%s", err)

	}
	fmt.Printf("更新数量 : %d\n", updateCount)

	//批量更新
	filter = mongodb.Newfilter().IN("year_published", 11, 12, 13)
	updateCount, err = mgdb.UpdateMany(ctx, filter, map[string]interface{}{
		"author": "update author",
	})

	if err != nil {
		log.Fatalf("%s", err)
	}
	fmt.Printf("批量更新数量 : %d\n", updateCount)

	//单条数据删除
	filter = mongodb.Newfilter().EQ("year_published", 15)
	deleteCount, err := mgdb.DeleteOne(ctx, filter)
	if err != nil {
		log.Fatalf("%s", err)
	}
	fmt.Printf("单条数据删除数量 : %d\n", deleteCount)

	//多条数据删除
	filter = mongodb.Newfilter().IN("year_published", 16, 17, 18)
	deleteCount, err = mgdb.DeleteMany(ctx, filter)
	if err != nil {
		log.Fatalf("%s", err)
	}
	fmt.Printf("多条数据删除数量 : %d\n", deleteCount)

}