github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/kvstore/db.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  	"github.com/cockroachdb/pebble"
     8  	"sync"
     9  )
    10  
    11  type DB struct {
    12  	store *Store
    13  
    14  	name      string
    15  	namespace string
    16  	tableMap  map[string]TableInterface
    17  
    18  	mu sync.RWMutex
    19  }
    20  
    21  func NewDB(store *Store, dbName string) (*DB, error) {
    22  	if !IsValidName(dbName) {
    23  		return nil, errors.New("invalid database name '" + dbName + "'")
    24  	}
    25  
    26  	return &DB{
    27  		store:     store,
    28  		name:      dbName,
    29  		namespace: "$" + dbName + "$",
    30  		tableMap:  map[string]TableInterface{},
    31  	}, nil
    32  }
    33  
    34  func (this *DB) AddTable(table TableInterface) {
    35  	table.SetNamespace([]byte(this.Namespace() + table.Name() + "$"))
    36  	table.SetDB(this)
    37  
    38  	this.mu.Lock()
    39  	defer this.mu.Unlock()
    40  
    41  	this.tableMap[table.Name()] = table
    42  }
    43  
    44  func (this *DB) Name() string {
    45  	return this.name
    46  }
    47  
    48  func (this *DB) Namespace() string {
    49  	return this.namespace
    50  }
    51  
    52  func (this *DB) Store() *Store {
    53  	return this.store
    54  }
    55  
    56  func (this *DB) Inspect(fn func(key []byte, value []byte)) error {
    57  	it, err := this.store.rawDB.NewIter(&pebble.IterOptions{
    58  		LowerBound: []byte(this.namespace),
    59  		UpperBound: append([]byte(this.namespace), 0xFF, 0xFF),
    60  	})
    61  	if err != nil {
    62  		return err
    63  	}
    64  	defer func() {
    65  		_ = it.Close()
    66  	}()
    67  
    68  	for it.First(); it.Valid(); it.Next() {
    69  		value, valueErr := it.ValueAndErr()
    70  		if valueErr != nil {
    71  			return valueErr
    72  		}
    73  		fn(it.Key(), value)
    74  	}
    75  
    76  	return nil
    77  }
    78  
    79  // Truncate the database
    80  func (this *DB) Truncate() error {
    81  	this.mu.Lock()
    82  	defer this.mu.Unlock()
    83  
    84  	var start = []byte(this.Namespace())
    85  	return this.store.rawDB.DeleteRange(start, append(start, 0xFF), DefaultWriteOptions)
    86  }
    87  
    88  func (this *DB) Close() error {
    89  	this.mu.Lock()
    90  	defer this.mu.Unlock()
    91  
    92  	var lastErr error
    93  	for _, table := range this.tableMap {
    94  		err := table.Close()
    95  		if err != nil {
    96  			lastErr = err
    97  		}
    98  	}
    99  
   100  	return lastErr
   101  }