github.com/TeaOSLab/EdgeNode@v1.3.8/internal/utils/agents/db_kv.go (about)

     1  // Copyright 2024 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package agents
     4  
     5  import (
     6  	"errors"
     7  	"github.com/TeaOSLab/EdgeNode/internal/events"
     8  	"github.com/TeaOSLab/EdgeNode/internal/utils/kvstore"
     9  )
    10  
    11  type KVDB struct {
    12  	table   *kvstore.Table[*AgentIP]
    13  	encoder *AgentIPEncoder[*AgentIP]
    14  	lastKey string
    15  }
    16  
    17  func NewKVDB() *KVDB {
    18  	var db = &KVDB{}
    19  
    20  	events.OnClose(func() {
    21  		_ = db.Close()
    22  	})
    23  
    24  	return db
    25  }
    26  
    27  func (this *KVDB) Init() error {
    28  	store, err := kvstore.DefaultStore()
    29  	if err != nil {
    30  		return err
    31  	}
    32  
    33  	db, err := store.NewDB("agents")
    34  	if err != nil {
    35  		return err
    36  	}
    37  
    38  	{
    39  		this.encoder = &AgentIPEncoder[*AgentIP]{}
    40  		table, tableErr := kvstore.NewTable[*AgentIP]("agent_ips", this.encoder)
    41  		if tableErr != nil {
    42  			return tableErr
    43  		}
    44  		db.AddTable(table)
    45  		this.table = table
    46  	}
    47  
    48  	return nil
    49  }
    50  
    51  func (this *KVDB) InsertAgentIP(ipId int64, ip string, agentCode string) error {
    52  	if this.table == nil {
    53  		return errors.New("table should not be nil")
    54  	}
    55  
    56  	var item = &AgentIP{
    57  		Id:        ipId,
    58  		IP:        ip,
    59  		AgentCode: agentCode,
    60  	}
    61  	var key = this.encoder.EncodeKey(item)
    62  	return this.table.Set(key, item)
    63  }
    64  
    65  func (this *KVDB) ListAgentIPs(offset int64, size int64) (agentIPs []*AgentIP, err error) {
    66  	if this.table == nil {
    67  		return nil, errors.New("table should not be nil")
    68  	}
    69  
    70  	err = this.table.
    71  		Query().
    72  		Limit(int(size)).
    73  		Offset(this.lastKey).
    74  		FindAll(func(tx *kvstore.Tx[*AgentIP], item kvstore.Item[*AgentIP]) (goNext bool, err error) {
    75  			this.lastKey = item.Key
    76  			agentIPs = append(agentIPs, item.Value)
    77  			return true, nil
    78  		})
    79  
    80  	return
    81  }
    82  
    83  func (this *KVDB) Close() error {
    84  	return nil
    85  }
    86  
    87  func (this *KVDB) Flush() error {
    88  	if this.table == nil {
    89  		return errors.New("table should not be nil")
    90  	}
    91  
    92  	return this.table.DB().Store().Flush()
    93  }