package mongo import ( "context" "fmt" "github.com/zeromicro/go-zero/core/logx" "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("mongo: document db: %w", err) } return &DocumentDBWithCache{ DocumentDBUseCase: documentDB, Cache: MustModelCache(cacheConf, cacheOpts...), }, nil } func (dc *DocumentDBWithCache) DelCache(ctx context.Context, keys ...string) error { return dc.Cache.DelCtx(ctx, keys...) } func (dc *DocumentDBWithCache) GetCache(ctx context.Context, key string, v any) error { return dc.Cache.GetCtx(ctx, key, v) } func (dc *DocumentDBWithCache) SetCache(ctx context.Context, key string, v any) error { return dc.Cache.SetCtx(ctx, key, v) } func (dc *DocumentDBWithCache) DeleteOne(ctx context.Context, key string, filter any, opts ...options.Lister[options.DeleteOneOptions]) (int64, error) { val, err := dc.GetClient().DeleteOne(ctx, filter, opts...) if err != nil { return 0, err } dc.delCacheBestEffort(ctx, key) return val, nil } func (dc *DocumentDBWithCache) FindOne(ctx context.Context, key string, v, filter any, opts ...options.Lister[options.FindOneOptions]) error { return dc.Cache.TakeCtx(ctx, v, key, func(v any) error { return dc.GetClient().FindOne(ctx, v, filter, opts...) }) } func (dc *DocumentDBWithCache) FindOneAndDelete(ctx context.Context, key string, v, filter any, opts ...options.Lister[options.FindOneAndDeleteOptions]) error { if err := dc.GetClient().FindOneAndDelete(ctx, v, filter, opts...); err != nil { return err } dc.delCacheBestEffort(ctx, key) return nil } func (dc *DocumentDBWithCache) FindOneAndReplace(ctx context.Context, key string, v, filter, replacement any, opts ...options.Lister[options.FindOneAndReplaceOptions]) error { if err := dc.GetClient().FindOneAndReplace(ctx, v, filter, replacement, opts...); err != nil { return err } dc.delCacheBestEffort(ctx, key) return nil } func (dc *DocumentDBWithCache) InsertOne(ctx context.Context, key string, document any, opts ...options.Lister[options.InsertOneOptions]) (*mongo.InsertOneResult, error) { res, err := dc.GetClient().InsertOne(ctx, document, opts...) if err != nil { return nil, err } dc.delCacheBestEffort(ctx, key) return res, nil } func (dc *DocumentDBWithCache) UpdateByID(ctx context.Context, key string, id, update any, opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) { res, err := dc.GetClient().UpdateByID(ctx, id, update, opts...) if err != nil { return nil, err } dc.delCacheBestEffort(ctx, key) return res, nil } func (dc *DocumentDBWithCache) UpdateMany(ctx context.Context, keys []string, filter, update any, opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error) { res, err := dc.GetClient().UpdateMany(ctx, filter, update, opts...) if err != nil { return nil, err } dc.delCacheBestEffort(ctx, keys...) return res, nil } func (dc *DocumentDBWithCache) UpdateOne(ctx context.Context, key string, filter, update any, opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error) { res, err := dc.GetClient().UpdateOne(ctx, filter, update, opts...) if err != nil { return nil, err } dc.delCacheBestEffort(ctx, key) return res, nil } func (dc *DocumentDBWithCache) delCacheBestEffort(ctx context.Context, keys ...string) { if len(keys) == 0 { return } if err := dc.DelCache(ctx, keys...); err != nil { logx.WithContext(ctx).Errorf("[DocumentDBWithCache] del cache keys=%v: %v", keys, err) } } func MustModelCache(conf cache.CacheConf, opts ...cache.Option) cache.Cache { return cache.New(conf, singleFlight, stats, ErrNotFound, opts...) }