github.com/vicanso/pike@v1.0.1-0.20210630235453-9099e041f6ec/store/badger.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  	"fmt"
    27  	"strings"
    28  	"time"
    29  
    30  	badger "github.com/dgraph-io/badger/v3"
    31  	badgerOptions "github.com/dgraph-io/badger/v3/options"
    32  	"github.com/vicanso/pike/log"
    33  	"go.uber.org/zap"
    34  )
    35  
    36  type badgerStore struct {
    37  	db *badger.DB
    38  }
    39  
    40  type badgerLogger struct{}
    41  
    42  func (bl *badgerLogger) Errorf(format string, args ...interface{}) {
    43  	msg := strings.TrimSpace(fmt.Sprintf(format, args...))
    44  	log.Default().Error(msg,
    45  		zap.String("category", "badger"),
    46  	)
    47  }
    48  func (bl *badgerLogger) Warningf(format string, args ...interface{}) {
    49  	msg := strings.TrimSpace(fmt.Sprintf(format, args...))
    50  	log.Default().Warn(msg,
    51  		zap.String("category", "badger"),
    52  	)
    53  }
    54  func (bl *badgerLogger) Infof(format string, args ...interface{}) {
    55  	msg := strings.TrimSpace(fmt.Sprintf(format, args...))
    56  	log.Default().Info(msg,
    57  		zap.String("category", "badger"),
    58  	)
    59  }
    60  func (bl *badgerLogger) Debugf(format string, args ...interface{}) {
    61  	msg := strings.TrimSpace(fmt.Sprintf(format, args...))
    62  	log.Default().Debug(msg,
    63  		zap.String("category", "badger"),
    64  	)
    65  }
    66  
    67  // newBadgerStore create a new badger store
    68  func newBadgerStore(path string) (Store, error) {
    69  	options := badger.DefaultOptions(path)
    70  	// 为了更高的性能，数据不压缩
    71  	options.Compression = badgerOptions.None
    72  	options.Logger = &badgerLogger{}
    73  	db, err := badger.Open(options)
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  	return &badgerStore{
    78  		db: db,
    79  	}, nil
    80  }
    81  
    82  // Get get data from badger
    83  func (bs *badgerStore) Get(key []byte) (data []byte, err error) {
    84  	err = bs.db.View(func(txn *badger.Txn) error {
    85  		item, err := txn.Get(key)
    86  		if err != nil {
    87  			if err == badger.ErrKeyNotFound {
    88  				err = ErrNotFound
    89  			}
    90  			return err
    91  		}
    92  		return item.Value(func(val []byte) error {
    93  			data = append([]byte{}, val...)
    94  			return nil
    95  		})
    96  	})
    97  	if err != nil {
    98  		return
    99  	}
   100  	return
   101  }
   102  
   103  // Set set data to badger
   104  func (bs *badgerStore) Set(key []byte, data []byte, ttl time.Duration) (err error) {
   105  	return bs.db.Update(func(txn *badger.Txn) error {
   106  		e := badger.NewEntry(key, data).
   107  			WithTTL(ttl)
   108  		return txn.SetEntry(e)
   109  	})
   110  }
   111  
   112  // Delete delete data from badger
   113  func (bs *badgerStore) Delete(key []byte) (err error) {
   114  	return bs.db.Update(func(txn *badger.Txn) error {
   115  		return txn.Delete(key)
   116  	})
   117  }
   118  
   119  // Close close badger
   120  func (bs *badgerStore) Close() error {
   121  	return bs.db.Close()
   122  }