ariga.io/entcache@v0.1.1-0.20230620164151-0eb723a11c40/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  //
    45  func Skip(ctx context.Context) context.Context {
    46  	c, ok := ctx.Value(ctxOptionsKey).(*ctxOptions)
    47  	if !ok {
    48  		return context.WithValue(ctx, ctxOptionsKey, &ctxOptions{skip: true})
    49  	}
    50  	c.skip = true
    51  	return ctx
    52  }
    53  
    54  // Evict returns a new Context that tells the Driver
    55  // to skip and invalidate the cache entry on Query.
    56  //
    57  //	client.T.Query().All(entcache.Evict(ctx))
    58  //
    59  func Evict(ctx context.Context) context.Context {
    60  	c, ok := ctx.Value(ctxOptionsKey).(*ctxOptions)
    61  	if !ok {
    62  		return context.WithValue(ctx, ctxOptionsKey, &ctxOptions{skip: true, evict: true})
    63  	}
    64  	c.skip = true
    65  	c.evict = true
    66  	return ctx
    67  }
    68  
    69  // WithKey returns a new Context that carries the Key for the cache entry.
    70  // Note that, this option should not be used if the ent.Client query involves
    71  // more than 1 SQL query (e.g. eager loading).
    72  //
    73  //	client.T.Query().All(entcache.WithKey(ctx, "key"))
    74  //
    75  func WithKey(ctx context.Context, key Key) context.Context {
    76  	c, ok := ctx.Value(ctxOptionsKey).(*ctxOptions)
    77  	if !ok {
    78  		return context.WithValue(ctx, ctxOptionsKey, &ctxOptions{key: key})
    79  	}
    80  	c.key = key
    81  	return ctx
    82  }
    83  
    84  // WithTTL returns a new Context that carries the TTL for the cache entry.
    85  //
    86  //	client.T.Query().All(entcache.WithTTL(ctx, time.Second))
    87  //
    88  func WithTTL(ctx context.Context, ttl time.Duration) context.Context {
    89  	c, ok := ctx.Value(ctxOptionsKey).(*ctxOptions)
    90  	if !ok {
    91  		return context.WithValue(ctx, ctxOptionsKey, &ctxOptions{ttl: ttl})
    92  	}
    93  	c.ttl = ttl
    94  	return ctx
    95  }