github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/kvstore/tx.go (about)

     1  // Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package kvstore
     4  
     5  import (
     6  	"errors"
     7  	"fmt"
     8  	"github.com/cockroachdb/pebble"
     9  )
    10  
    11  type Tx[T any] struct {
    12  	table    *Table[T]
    13  	readOnly bool
    14  
    15  	batch *pebble.Batch
    16  }
    17  
    18  func NewTx[T any](table *Table[T], readOnly bool) (*Tx[T], error) {
    19  	if table.db == nil {
    20  		return nil, errors.New("the table has not been added to a db")
    21  	}
    22  	if table.db.store == nil {
    23  		return nil, errors.New("the db has not been added to a store")
    24  	}
    25  	if table.db.store.rawDB == nil {
    26  		return nil, errors.New("the store has not been opened")
    27  	}
    28  
    29  	return &Tx[T]{
    30  		table:    table,
    31  		readOnly: readOnly,
    32  		batch:    table.db.store.rawDB.NewIndexedBatch(),
    33  	}, nil
    34  }
    35  
    36  func (this *Tx[T]) Set(key string, value T) error {
    37  	if this.readOnly {
    38  		return errors.New("can not set value in readonly transaction")
    39  	}
    40  
    41  	if len(key) > KeyMaxLength {
    42  		return ErrKeyTooLong
    43  	}
    44  
    45  	valueBytes, err := this.table.encoder.Encode(value)
    46  	if err != nil {
    47  		return err
    48  	}
    49  
    50  	return this.table.set(this, key, valueBytes, value, false, false)
    51  }
    52  
    53  func (this *Tx[T]) SetSync(key string, value T) error {
    54  	if this.readOnly {
    55  		return errors.New("can not set value in readonly transaction")
    56  	}
    57  
    58  	if len(key) > KeyMaxLength {
    59  		return ErrKeyTooLong
    60  	}
    61  
    62  	valueBytes, err := this.table.encoder.Encode(value)
    63  	if err != nil {
    64  		return err
    65  	}
    66  
    67  	return this.table.set(this, key, valueBytes, value, false, true)
    68  }
    69  
    70  func (this *Tx[T]) Insert(key string, value T) error {
    71  	if this.readOnly {
    72  		return errors.New("can not set value in readonly transaction")
    73  	}
    74  
    75  	if len(key) > KeyMaxLength {
    76  		return ErrKeyTooLong
    77  	}
    78  
    79  	valueBytes, err := this.table.encoder.Encode(value)
    80  	if err != nil {
    81  		return err
    82  	}
    83  
    84  	return this.table.set(this, key, valueBytes, value, true, false)
    85  }
    86  
    87  func (this *Tx[T]) Get(key string) (value T, err error) {
    88  	if this.table.isClosed {
    89  		err = NewTableClosedErr(this.table.name)
    90  		return
    91  	}
    92  	return this.table.get(this, key)
    93  }
    94  
    95  func (this *Tx[T]) Delete(key string) error {
    96  	if this.table.isClosed {
    97  		return NewTableClosedErr(this.table.name)
    98  	}
    99  	if this.readOnly {
   100  		return errors.New("can not delete value in readonly transaction")
   101  	}
   102  
   103  	return this.table.deleteKeys(this, key)
   104  }
   105  
   106  func (this *Tx[T]) NewIterator(opt *IteratorOptions) (*pebble.Iterator, error) {
   107  	return this.batch.NewIter(opt.RawOptions())
   108  }
   109  
   110  func (this *Tx[T]) Close() error {
   111  	return this.batch.Close()
   112  }
   113  
   114  func (this *Tx[T]) Commit() (err error) {
   115  	return this.commit(DefaultWriteOptions)
   116  }
   117  
   118  func (this *Tx[T]) CommitSync() (err error) {
   119  	return this.commit(DefaultWriteSyncOptions)
   120  }
   121  
   122  func (this *Tx[T]) Query() *Query[T] {
   123  	var query = NewQuery[T]()
   124  	query.SetTx(this)
   125  	return query
   126  }
   127  
   128  func (this *Tx[T]) commit(opt *pebble.WriteOptions) (err error) {
   129  	defer func() {
   130  		var panicErr = recover()
   131  		if panicErr != nil {
   132  			resultErr, ok := panicErr.(error)
   133  			if ok {
   134  				err = fmt.Errorf("commit batch failed: %w", resultErr)
   135  			}
   136  		}
   137  	}()
   138  
   139  	return this.batch.Commit(opt)
   140  }