package model import ( "context" "errors" "time" "github.com/zeromicro/go-zero/core/stores/monc" "go.mongodb.org/mongo-driver/mongo/options" "github.com/zeromicro/go-zero/core/logx" "github.com/zeromicro/go-zero/core/stores/mon" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/primitive" "go.mongodb.org/mongo-driver/mongo" ) var _ PostModel = (*customPostModel)(nil) type ( // PostModel is an interface to be customized, add more methods here, // and implement the added methods in customPostModel. PostModel interface { postModel DeleteMany(ctx context.Context, id ...string) (int64, error) UpdateOptional(ctx context.Context, data *Post) (*mongo.UpdateResult, error) Find(ctx context.Context, param *QueryPostModelReq) ([]*Post, int64, error) } customPostModel struct { *defaultPostModel } QueryPostModelReq struct { UID []string Id []string OnlyAds *bool PageSize int64 PageIndex int64 } ) // NewPostModel returns a model for the mongo. func NewPostModel(url, db, collection string) PostModel { conn := mon.MustNewModel(url, db, collection) return &customPostModel{ defaultPostModel: newDefaultPostModel(conn), } } func (m *defaultPostModel) DeleteMany(ctx context.Context, id ...string) (int64, error) { objectIDs := make([]primitive.ObjectID, 0, len(id)) // prepare for _, item := range id { oid, err := primitive.ObjectIDFromHex(item) if err != nil { logx.WithCallerSkip(1).WithFields( logx.Field("func", "defaultPostModel.DeleteMany"), logx.Field("id", item), ).Error(err.Error()) continue } objectIDs = append(objectIDs, oid) } // 檢查是否有有效的 ObjectIDs if len(objectIDs) == 0 { return 0, ErrNotFound } // 刪除文檔 res, err := m.conn.DeleteMany(ctx, bson.M{"_id": bson.M{"$in": objectIDs}}) if err != nil { return 0, err } return res, err } func (m *defaultPostModel) UpdateOptional(ctx context.Context, data *Post) (*mongo.UpdateResult, error) { update := bson.M{"$set": bson.M{}} if data.UID != "" { update["$set"].(bson.M)["uid"] = data.UID } if data.Content != "" { update["$set"].(bson.M)["content"] = data.Content } if data.Status != 0 { update["$set"].(bson.M)["status"] = data.Status } if data.IsAd { update["$set"].(bson.M)["is_ad"] = data.IsAd } if len(data.Tags) > 0 { update["$set"].(bson.M)["tags"] = data.Tags } if len(data.MediaURL) > 0 { update["$set"].(bson.M)["media_url"] = data.MediaURL } if data.Like != -1 { update["$set"].(bson.M)["like"] = data.Like } if data.DisLike != -1 { update["$set"].(bson.M)["dislike"] = data.DisLike } // UpdateAt 是每次都需要更新的,不用檢查 update["$set"].(bson.M)["updateAt"] = time.Now().UTC().UnixNano() res, err := m.conn.UpdateOne(ctx, bson.M{"_id": data.ID}, update) return res, err } // Find 貼文列表 func (m *defaultPostModel) Find(ctx context.Context, param *QueryPostModelReq) ([]*Post, int64, error) { filter := bson.M{} // 添加 UID 過濾條件 if len(param.UID) > 0 { filter["uid"] = bson.M{"$in": param.UID} } // 添加 ID 過濾條件 if len(param.Id) > 0 { var ids []primitive.ObjectID for _, item := range param.Id { oid, err := primitive.ObjectIDFromHex(item) if err != nil { // log continue } ids = append(ids, oid) } filter["_id"] = bson.M{"$in": ids} } // 添加 OnlyAds 過濾條件 if param.OnlyAds != nil { // true 是廣告 false 不是廣告 , 沒寫就是不過率 filter["is_ad"] = *param.OnlyAds } // 分頁處理 opts := options.Find() opts.SetSort(bson.D{{"create_time", -1}}) if param.PageSize > 0 { opts.SetLimit(param.PageSize) } if param.PageIndex > 0 { opts.SetSkip((param.PageIndex - 1) * param.PageSize) } // 計算總數(不考慮分頁) totalCount, err := m.conn.CountDocuments(ctx, filter) if err != nil { return nil, 0, err } result := make([]*Post, 0, param.PageSize) // 執行查詢 err = m.conn.Find(ctx, &result, filter, opts) if err != nil { return nil, 0, err } switch { case err == nil: return result, totalCount, nil case errors.Is(err, monc.ErrNotFound): return nil, 0, ErrNotFound default: return nil, 0, err } }