github.com/shuguocloud/go-zero@v1.3.0/core/stores/sqlc/cachedsql.go (about)

     1  package sqlc
     2  
     3  import (
     4  	"database/sql"
     5  	"time"
     6  
     7  	"github.com/shuguocloud/go-zero/core/stores/cache"
     8  	"github.com/shuguocloud/go-zero/core/stores/redis"
     9  	"github.com/shuguocloud/go-zero/core/stores/sqlx"
    10  	"github.com/shuguocloud/go-zero/core/syncx"
    11  )
    12  
    13  // see doc/sql-cache.md
    14  const cacheSafeGapBetweenIndexAndPrimary = time.Second * 5
    15  
    16  var (
    17  	// ErrNotFound is an alias of sqlx.ErrNotFound.
    18  	ErrNotFound = sqlx.ErrNotFound
    19  
    20  	// can't use one SingleFlight per conn, because multiple conns may share the same cache key.
    21  	exclusiveCalls = syncx.NewSingleFlight()
    22  	stats          = cache.NewStat("sqlc")
    23  )
    24  
    25  type (
    26  	// ExecFn defines the sql exec method.
    27  	ExecFn func(conn sqlx.SqlConn) (sql.Result, error)
    28  	// IndexQueryFn defines the query method that based on unique indexes.
    29  	IndexQueryFn func(conn sqlx.SqlConn, v interface{}) (interface{}, error)
    30  	// PrimaryQueryFn defines the query method that based on primary keys.
    31  	PrimaryQueryFn func(conn sqlx.SqlConn, v, primary interface{}) error
    32  	// QueryFn defines the query method.
    33  	QueryFn func(conn sqlx.SqlConn, v interface{}) error
    34  
    35  	// A CachedConn is a DB connection with cache capability.
    36  	CachedConn struct {
    37  		db    sqlx.SqlConn
    38  		cache cache.Cache
    39  	}
    40  )
    41  
    42  // NewConn returns a CachedConn with a redis cluster cache.
    43  func NewConn(db sqlx.SqlConn, c cache.CacheConf, opts ...cache.Option) CachedConn {
    44  	cc := cache.New(c, exclusiveCalls, stats, sql.ErrNoRows, opts...)
    45  	return NewConnWithCache(db, cc)
    46  }
    47  
    48  // NewConnWithCache returns a CachedConn with a custom cache.
    49  func NewConnWithCache(db sqlx.SqlConn, c cache.Cache) CachedConn {
    50  	return CachedConn{
    51  		db:    db,
    52  		cache: c,
    53  	}
    54  }
    55  
    56  // NewNodeConn returns a CachedConn with a redis node cache.
    57  func NewNodeConn(db sqlx.SqlConn, rds *redis.Redis, opts ...cache.Option) CachedConn {
    58  	c := cache.NewNode(rds, exclusiveCalls, stats, sql.ErrNoRows, opts...)
    59  	return NewConnWithCache(db, c)
    60  }
    61  
    62  // DelCache deletes cache with keys.
    63  func (cc CachedConn) DelCache(keys ...string) error {
    64  	return cc.cache.Del(keys...)
    65  }
    66  
    67  // GetCache unmarshals cache with given key into v.
    68  func (cc CachedConn) GetCache(key string, v interface{}) error {
    69  	return cc.cache.Get(key, v)
    70  }
    71  
    72  // Exec runs given exec on given keys, and returns execution result.
    73  func (cc CachedConn) Exec(exec ExecFn, keys ...string) (sql.Result, error) {
    74  	res, err := exec(cc.db)
    75  	if err != nil {
    76  		return nil, err
    77  	}
    78  
    79  	if err := cc.DelCache(keys...); err != nil {
    80  		return nil, err
    81  	}
    82  
    83  	return res, nil
    84  }
    85  
    86  // ExecNoCache runs exec with given sql statement, without affecting cache.
    87  func (cc CachedConn) ExecNoCache(q string, args ...interface{}) (sql.Result, error) {
    88  	return cc.db.Exec(q, args...)
    89  }
    90  
    91  // QueryRow unmarshals into v with given key and query func.
    92  func (cc CachedConn) QueryRow(v interface{}, key string, query QueryFn) error {
    93  	return cc.cache.Take(v, key, func(v interface{}) error {
    94  		return query(cc.db, v)
    95  	})
    96  }
    97  
    98  // QueryRowIndex unmarshals into v with given key.
    99  func (cc CachedConn) QueryRowIndex(v interface{}, key string, keyer func(primary interface{}) string,
   100  	indexQuery IndexQueryFn, primaryQuery PrimaryQueryFn) error {
   101  	var primaryKey interface{}
   102  	var found bool
   103  
   104  	if err := cc.cache.TakeWithExpire(&primaryKey, key, func(val interface{}, expire time.Duration) (err error) {
   105  		primaryKey, err = indexQuery(cc.db, v)
   106  		if err != nil {
   107  			return
   108  		}
   109  
   110  		found = true
   111  		return cc.cache.SetWithExpire(keyer(primaryKey), v, expire+cacheSafeGapBetweenIndexAndPrimary)
   112  	}); err != nil {
   113  		return err
   114  	}
   115  
   116  	if found {
   117  		return nil
   118  	}
   119  
   120  	return cc.cache.Take(v, keyer(primaryKey), func(v interface{}) error {
   121  		return primaryQuery(cc.db, v, primaryKey)
   122  	})
   123  }
   124  
   125  // QueryRowNoCache unmarshals into v with given statement.
   126  func (cc CachedConn) QueryRowNoCache(v interface{}, q string, args ...interface{}) error {
   127  	return cc.db.QueryRow(v, q, args...)
   128  }
   129  
   130  // QueryRowsNoCache unmarshals into v with given statement.
   131  // It doesn't use cache, because it might cause consistency problem.
   132  func (cc CachedConn) QueryRowsNoCache(v interface{}, q string, args ...interface{}) error {
   133  	return cc.db.QueryRows(v, q, args...)
   134  }
   135  
   136  // SetCache sets v into cache with given key.
   137  func (cc CachedConn) SetCache(key string, v interface{}) error {
   138  	return cc.cache.Set(key, v)
   139  }
   140  
   141  // Transact runs given fn in transaction mode.
   142  func (cc CachedConn) Transact(fn func(sqlx.Session) error) error {
   143  	return cc.db.Transact(fn)
   144  }