github.com/DeltaLaboratory/entcache@v0.1.1/context.go (about)

     1  package entcache
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  )
     7  
     8  type ctxKey struct{}
     9  
    10  // NewContext returns a new Context that carries a cache.
    11  func NewContext(ctx context.Context, levels ...AddGetDeleter) context.Context {
    12  	var cache AddGetDeleter
    13  	switch len(levels) {
    14  	case 0:
    15  		cache = NewLRU(0)
    16  	case 1:
    17  		cache = levels[0]
    18  	default:
    19  		cache = &multiLevel{levels: levels}
    20  	}
    21  	return context.WithValue(ctx, ctxKey{}, cache)
    22  }
    23  
    24  // FromContext returns the cache value stored in ctx, if any.
    25  func FromContext(ctx context.Context) (AddGetDeleter, bool) {
    26  	c, ok := ctx.Value(ctxKey{}).(AddGetDeleter)
    27  	return c, ok
    28  }
    29  
    30  // ctxOptions allows injecting runtime options.
    31  type ctxOptions struct {
    32  	skip  bool          // i.e. skip entry.
    33  	evict bool          // i.e. skip and invalidate entry.
    34  	key   Key           // entry key.
    35  	ttl   time.Duration // entry duration.
    36  }
    37  
    38  var ctxOptionsKey ctxOptions
    39  
    40  // Skip returns a new Context that tells the Driver
    41  // to skip the cache entry on Query.
    42  //
    43  //	client.T.Query().All(entcache.Skip(ctx))
    44  func Skip(ctx context.Context) context.Context {
    45  	c, ok := ctx.Value(ctxOptionsKey).(*ctxOptions)
    46  	if !ok {
    47  		return context.WithValue(ctx, ctxOptionsKey, &ctxOptions{skip: true})
    48  	}
    49  	c.skip = true
    50  	return ctx
    51  }
    52  
    53  // Evict returns a new Context that tells the Driver
    54  // to skip and invalidate the cache entry on Query.
    55  //
    56  //	client.T.Query().All(entcache.Evict(ctx))
    57  func Evict(ctx context.Context) context.Context {
    58  	c, ok := ctx.Value(ctxOptionsKey).(*ctxOptions)
    59  	if !ok {
    60  		return context.WithValue(ctx, ctxOptionsKey, &ctxOptions{skip: true, evict: true})
    61  	}
    62  	c.skip = true
    63  	c.evict = true
    64  	return ctx
    65  }
    66  
    67  // WithKey returns a new Context that carries the Key for the cache entry.
    68  // Note that this option should not be used if the ent.Client query involves
    69  // more than 1 SQL query (e.g., eager loading).
    70  //
    71  //	client.T.Query().All(entcache.WithKey(ctx, "key"))
    72  func WithKey(ctx context.Context, key Key) context.Context {
    73  	c, ok := ctx.Value(ctxOptionsKey).(*ctxOptions)
    74  	if !ok {
    75  		return context.WithValue(ctx, ctxOptionsKey, &ctxOptions{key: key})
    76  	}
    77  	c.key = key
    78  	return ctx
    79  }
    80  
    81  // WithTTL returns a new Context that carries the TTL for the cache entry.
    82  //
    83  //	client.T.Query().All(entcache.WithTTL(ctx, time.Second))
    84  func WithTTL(ctx context.Context, ttl time.Duration) context.Context {
    85  	c, ok := ctx.Value(ctxOptionsKey).(*ctxOptions)
    86  	if !ok {
    87  		return context.WithValue(ctx, ctxOptionsKey, &ctxOptions{ttl: ttl})
    88  	}
    89  	c.ttl = ttl
    90  	return ctx
    91  }