github.com/masterhung0112/hk_server/v5@v5.0.0-20220302090640-ec71aef15e1c/store/sqlstore/context.go (about)

     1  package sqlstore
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/mattermost/gorp"
     7  )
     8  
     9  // storeContextKey is the base type for all context keys for the store.
    10  type storeContextKey string
    11  
    12  // contextValue is a type to hold some pre-determined context values.
    13  type contextValue string
    14  
    15  // Different possible values of contextValue.
    16  const (
    17  	useMaster contextValue = "useMaster"
    18  )
    19  
    20  // WithMaster adds the context value that master DB should be selected for this request.
    21  func WithMaster(ctx context.Context) context.Context {
    22  	return context.WithValue(ctx, storeContextKey(useMaster), true)
    23  }
    24  
    25  // hasMaster is a helper function to check whether master DB should be selected or not.
    26  func hasMaster(ctx context.Context) bool {
    27  	if v := ctx.Value(storeContextKey(useMaster)); v != nil {
    28  		if res, ok := v.(bool); ok && res {
    29  			return true
    30  		}
    31  	}
    32  	return false
    33  }
    34  
    35  // DBFromContext is a helper utility that returns the DB handle from a given context.
    36  func (ss *SqlStore) DBFromContext(ctx context.Context) *gorp.DbMap {
    37  	if hasMaster(ctx) {
    38  		return ss.GetMaster()
    39  	}
    40  	return ss.GetReplica()
    41  }