github.com/vicanso/pike@v1.0.1-0.20210630235453-9099e041f6ec/store/redis.go (about)

     1  // MIT License
     2  
     3  // Copyright (c) 2021 Tree Xie
     4  
     5  // Permission is hereby granted, free of charge, to any person obtaining a copy
     6  // of this software and associated documentation files (the "Software"), to deal
     7  // in the Software without restriction, including without limitation the rights
     8  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     9  // copies of the Software, and to permit persons to whom the Software is
    10  // furnished to do so, subject to the following conditions:
    11  
    12  // The above copyright notice and this permission notice shall be included in all
    13  // copies or substantial portions of the Software.
    14  
    15  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    16  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    17  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    18  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    19  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    20  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    21  // SOFTWARE.
    22  
    23  package store
    24  
    25  import (
    26  	"context"
    27  	"fmt"
    28  	"net/url"
    29  	"strconv"
    30  	"strings"
    31  	"time"
    32  
    33  	"github.com/go-redis/redis/v8"
    34  	"github.com/vicanso/pike/log"
    35  	"go.uber.org/zap"
    36  )
    37  
    38  type redisStore struct {
    39  	client redis.UniversalClient
    40  	// timeout 超时设置
    41  	timeout time.Duration
    42  	// prefix key的前缀
    43  	prefix string
    44  }
    45  
    46  type redisLogger struct{}
    47  
    48  func (rl *redisLogger) Printf(ctx context.Context, format string, v ...interface{}) {
    49  	log.Default().Info(fmt.Sprintf(format, v...),
    50  		zap.String("category", "redisLogger"),
    51  	)
    52  }
    53  
    54  func init() {
    55  	redis.SetLogger(&redisLogger{})
    56  }
    57  
    58  func newRedisStore(connectionURI string) (store Store, err error) {
    59  	urlInfo, err := url.Parse(connectionURI)
    60  	if err != nil {
    61  		return
    62  	}
    63  	user := ""
    64  	password := ""
    65  	if urlInfo.User != nil {
    66  		user = urlInfo.User.Username()
    67  		password, _ = urlInfo.User.Password()
    68  	}
    69  	// redis选择的db
    70  	db, _ := strconv.Atoi(urlInfo.Query().Get("db"))
    71  	// 设置的超时，如 3s
    72  	timeout, _ := time.ParseDuration(urlInfo.Query().Get("timeout"))
    73  	// 保存的key的前缀
    74  	prefix := urlInfo.Query().Get("prefix")
    75  
    76  	addrs := strings.Split(urlInfo.Host, ",")
    77  	master := urlInfo.Query().Get("master")
    78  
    79  	client := redis.NewUniversalClient(&redis.UniversalOptions{
    80  		Addrs:            addrs,
    81  		Username:         user,
    82  		Password:         password,
    83  		DB:               db,
    84  		SentinelPassword: password,
    85  		MasterName:       master,
    86  	})
    87  
    88  	// 默认3秒超时
    89  	if timeout == 0 {
    90  		timeout = 3 * time.Second
    91  	}
    92  	store = &redisStore{
    93  		client:  client,
    94  		timeout: timeout,
    95  		prefix:  prefix,
    96  	}
    97  	return
    98  }
    99  
   100  func (rs *redisStore) getKey(key []byte) string {
   101  	return rs.prefix + string(key)
   102  }
   103  
   104  // Get get data from redis
   105  func (rs *redisStore) Get(key []byte) (data []byte, err error) {
   106  	ctx, cancel := context.WithTimeout(context.Background(), rs.timeout)
   107  	defer cancel()
   108  	k := rs.getKey(key)
   109  	cmd := rs.client.Get(ctx, k)
   110  	data, err = cmd.Bytes()
   111  	if err != nil {
   112  		if err == redis.Nil {
   113  			err = ErrNotFound
   114  		}
   115  		return
   116  	}
   117  	return
   118  }
   119  
   120  // Set set data to redis
   121  func (rs *redisStore) Set(key []byte, data []byte, ttl time.Duration) (err error) {
   122  	ctx, cancel := context.WithTimeout(context.Background(), rs.timeout)
   123  	defer cancel()
   124  	k := rs.getKey(key)
   125  	cmd := rs.client.Set(ctx, k, data, ttl)
   126  	return cmd.Err()
   127  }
   128  
   129  // Delete delete date from redis
   130  func (rs *redisStore) Delete(key []byte) (err error) {
   131  	ctx, cancel := context.WithTimeout(context.Background(), rs.timeout)
   132  	defer cancel()
   133  	k := rs.getKey(key)
   134  	cmd := rs.client.Del(ctx, k)
   135  	return cmd.Err()
   136  }
   137  
   138  // Close close redis
   139  func (rs *redisStore) Close() error {
   140  	return rs.client.Close()
   141  }