package mongo import ( "context" "fmt" "github.com/zeromicro/go-zero/core/stores/cache" "github.com/zeromicro/go-zero/core/stores/mon" "go.mongodb.org/mongo-driver/v2/mongo" "go.mongodb.org/mongo-driver/v2/mongo/options" ) type DocumentDBWithCache struct { DocumentDBUseCase Cache cache.Cache } func MustDocumentDBWithCache(conf *Conf, collection string, cacheConf cache.CacheConf, dbOpts []mon.Option, cacheOpts []cache.Option) (DocumentDBWithCacheUseCase, error) { documentDB, err := NewDocumentDB(conf, collection, dbOpts...) if err != nil { return nil, fmt.Errorf("failed to initialize DocumentDB: %w", err) } c := MustModelCache(cacheConf, cacheOpts...) return &DocumentDBWithCache{ DocumentDBUseCase: documentDB, Cache: c, }, nil } func (dc *DocumentDBWithCache) DelCache(ctx context.Context, keys ...string) error { return dc.Cache.DelCtx(ctx, keys...) } func (dc *DocumentDBWithCache) GetCache(key string, v any) error { return dc.Cache.Get(key, v) } func (dc *DocumentDBWithCache) SetCache(key string, v any) error { return dc.Cache.Set(key, v) } // DeleteOne deletes a single document and invalidates cache func (dc *DocumentDBWithCache) DeleteOne(ctx context.Context, key string, filter any, opts ...*options.DeleteOneOptions) (int64, error) { // Convert options to Builder format var listerOpts []options.Lister[options.DeleteOneOptions] for _, opt := range opts { if opt != nil { builder := options.DeleteOne() if opt.Collation != nil { builder.SetCollation(opt.Collation) } if opt.Comment != nil { builder.SetComment(opt.Comment) } if opt.Hint != nil { builder.SetHint(opt.Hint) } listerOpts = append(listerOpts, builder) } } val, err := dc.GetClient().DeleteOne(ctx, filter, listerOpts...) if err != nil { return 0, err } if err := dc.DelCache(ctx, key); err != nil { return 0, err } return val, nil } // FindOne finds a single document with cache support func (dc *DocumentDBWithCache) FindOne(ctx context.Context, key string, v, filter any, opts ...*options.FindOneOptions) error { return dc.Cache.TakeCtx(ctx, v, key, func(v any) error { // Convert options to Builder format var listerOpts []options.Lister[options.FindOneOptions] for _, opt := range opts { if opt != nil { builder := options.FindOne() if opt.Collation != nil { builder.SetCollation(opt.Collation) } if opt.Comment != nil { builder.SetComment(opt.Comment) } if opt.Hint != nil { builder.SetHint(opt.Hint) } if opt.Projection != nil { builder.SetProjection(opt.Projection) } if opt.Skip != nil { builder.SetSkip(*opt.Skip) } if opt.Sort != nil { builder.SetSort(opt.Sort) } listerOpts = append(listerOpts, builder) } } return dc.GetClient().FindOne(ctx, v, filter, listerOpts...) }) } // FindOneAndDelete finds and deletes a single document with cache invalidation func (dc *DocumentDBWithCache) FindOneAndDelete(ctx context.Context, key string, v, filter any, opts ...*options.FindOneAndDeleteOptions) error { // Convert options to Builder format var listerOpts []options.Lister[options.FindOneAndDeleteOptions] for _, opt := range opts { if opt != nil { builder := options.FindOneAndDelete() if opt.Collation != nil { builder.SetCollation(opt.Collation) } if opt.Comment != nil { builder.SetComment(opt.Comment) } if opt.Hint != nil { builder.SetHint(opt.Hint) } if opt.Projection != nil { builder.SetProjection(opt.Projection) } if opt.Sort != nil { builder.SetSort(opt.Sort) } listerOpts = append(listerOpts, builder) } } if err := dc.GetClient().FindOneAndDelete(ctx, v, filter, listerOpts...); err != nil { return err } return dc.DelCache(ctx, key) } // FindOneAndReplace finds and replaces a single document with cache invalidation func (dc *DocumentDBWithCache) FindOneAndReplace(ctx context.Context, key string, v, filter, replacement any, opts ...*options.FindOneAndReplaceOptions) error { // Convert options to Builder format var listerOpts []options.Lister[options.FindOneAndReplaceOptions] for _, opt := range opts { if opt != nil { builder := options.FindOneAndReplace() if opt.Collation != nil { builder.SetCollation(opt.Collation) } if opt.Comment != nil { builder.SetComment(opt.Comment) } if opt.Hint != nil { builder.SetHint(opt.Hint) } if opt.Projection != nil { builder.SetProjection(opt.Projection) } if opt.ReturnDocument != nil { builder.SetReturnDocument(*opt.ReturnDocument) } if opt.Sort != nil { builder.SetSort(opt.Sort) } if opt.Upsert != nil { builder.SetUpsert(*opt.Upsert) } listerOpts = append(listerOpts, builder) } } if err := dc.GetClient().FindOneAndReplace(ctx, v, filter, replacement, listerOpts...); err != nil { return err } return dc.DelCache(ctx, key) } // InsertOne inserts a single document and invalidates cache func (dc *DocumentDBWithCache) InsertOne(ctx context.Context, key string, document any, opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) { // Convert options to Builder format var listerOpts []options.Lister[options.InsertOneOptions] for _, opt := range opts { if opt != nil { builder := options.InsertOne() if opt.BypassDocumentValidation != nil { builder.SetBypassDocumentValidation(*opt.BypassDocumentValidation) } if opt.Comment != nil { builder.SetComment(opt.Comment) } listerOpts = append(listerOpts, builder) } } res, err := dc.GetClient().Collection.InsertOne(ctx, document, listerOpts...) if err != nil { return nil, err } if err = dc.DelCache(ctx, key); err != nil { return nil, err } return res, nil } // UpdateByID updates a document by ID and invalidates cache func (dc *DocumentDBWithCache) UpdateByID(ctx context.Context, key string, id, update any, opts ...*options.UpdateOneOptions) (*mongo.UpdateResult, error) { // Convert options to Builder format var listerOpts []options.Lister[options.UpdateOneOptions] for _, opt := range opts { if opt != nil { builder := options.UpdateOne() if opt.ArrayFilters != nil { builder.SetArrayFilters(opt.ArrayFilters) } if opt.BypassDocumentValidation != nil { builder.SetBypassDocumentValidation(*opt.BypassDocumentValidation) } if opt.Collation != nil { builder.SetCollation(opt.Collation) } if opt.Comment != nil { builder.SetComment(opt.Comment) } if opt.Hint != nil { builder.SetHint(opt.Hint) } if opt.Upsert != nil { builder.SetUpsert(*opt.Upsert) } listerOpts = append(listerOpts, builder) } } res, err := dc.GetClient().Collection.UpdateByID(ctx, id, update, listerOpts...) if err != nil { return nil, err } if err = dc.DelCache(ctx, key); err != nil { return nil, err } return res, nil } // UpdateMany updates multiple documents and invalidates cache func (dc *DocumentDBWithCache) UpdateMany(ctx context.Context, keys []string, filter, update any, opts ...*options.UpdateManyOptions) (*mongo.UpdateResult, error) { // Convert options to Builder format var listerOpts []options.Lister[options.UpdateManyOptions] for _, opt := range opts { if opt != nil { builder := options.UpdateMany() if opt.ArrayFilters != nil { builder.SetArrayFilters(opt.ArrayFilters) } if opt.BypassDocumentValidation != nil { builder.SetBypassDocumentValidation(*opt.BypassDocumentValidation) } if opt.Collation != nil { builder.SetCollation(opt.Collation) } if opt.Comment != nil { builder.SetComment(opt.Comment) } if opt.Hint != nil { builder.SetHint(opt.Hint) } if opt.Upsert != nil { builder.SetUpsert(*opt.Upsert) } listerOpts = append(listerOpts, builder) } } res, err := dc.GetClient().Collection.UpdateMany(ctx, filter, update, listerOpts...) if err != nil { return nil, err } if err = dc.DelCache(ctx, keys...); err != nil { return nil, err } return res, nil } // UpdateOne updates a single document and invalidates cache func (dc *DocumentDBWithCache) UpdateOne(ctx context.Context, key string, filter, update any, opts ...*options.UpdateOneOptions) (*mongo.UpdateResult, error) { // Convert options to Builder format var listerOpts []options.Lister[options.UpdateOneOptions] for _, opt := range opts { if opt != nil { builder := options.UpdateOne() if opt.ArrayFilters != nil { builder.SetArrayFilters(opt.ArrayFilters) } if opt.BypassDocumentValidation != nil { builder.SetBypassDocumentValidation(*opt.BypassDocumentValidation) } if opt.Collation != nil { builder.SetCollation(opt.Collation) } if opt.Comment != nil { builder.SetComment(opt.Comment) } if opt.Hint != nil { builder.SetHint(opt.Hint) } if opt.Upsert != nil { builder.SetUpsert(*opt.Upsert) } listerOpts = append(listerOpts, builder) } } res, err := dc.GetClient().Collection.UpdateOne(ctx, filter, update, listerOpts...) if err != nil { return nil, err } if err = dc.DelCache(ctx, key); err != nil { return nil, err } return res, nil } // ======================== // MustModelCache returns a cache cluster. func MustModelCache(conf cache.CacheConf, opts ...cache.Option) cache.Cache { return cache.New(conf, singleFlight, stats, mongo.ErrNoDocuments, opts...) }