package localserver import ( "encoding/binary" "encoding/json" "os" "path/filepath" "sync" "github.com/dgraph-io/badger/v3" ) type NodeChatDBHandler struct { ChatID string PreviousId uint64 ItemCount uint64 db func(func(*badger.DB) (err error)) (err error) lock *sync.RWMutex } func NewNodeChatDBHandler(chatID string) (nodeChatDBHandler *NodeChatDBHandler, err error) { nodeChatDBHandler = &NodeChatDBHandler{ db: func(f func(*badger.DB) (err error)) (err error) { db, err := badger.Open(badger.DefaultOptions(filepath.Join(dataPath, "data", "chats", chatID, "messages")).WithLogger(dbLogger)) if err != nil { return } defer db.Close() err = f(db) return }, ChatID: chatID, lock: new(sync.RWMutex), } err = nodeChatDBHandler.db(func(d *badger.DB) (err error) { err = d.View(func(txn *badger.Txn) error { opt := badger.DefaultIteratorOptions it := txn.NewIterator(opt) defer it.Close() var counter uint64 = 1 var itemCounter uint64 = 0 for it.Rewind(); it.Valid(); it.Next() { item := it.Item() if err = item.Value(func(val []byte) (err error) { var chatMessage *ChatMessage if err = json.Unmarshal(val, &chatMessage); err != nil { return err } counter = chatMessage.ID itemCounter++ return }); err != nil { return err } } nodeChatDBHandler.PreviousId = counter nodeChatDBHandler.ItemCount = itemCounter return nil }) return }) return } func (ncdbh *NodeChatDBHandler) calculateNewChatCount(previousId uint64) (count uint, err error) { err = ncdbh.db(func(d *badger.DB) (err error) { err = d.View(func(txn *badger.Txn) error { opt := badger.DefaultIteratorOptions it := txn.NewIterator(opt) defer it.Close() count = 0 b := make([]byte, bufferSize) binary.BigEndian.PutUint64(b, previousId) if previousId == 0 || previousId == 1 { count++ for it.Rewind(); it.Valid(); it.Next() { count++ } } else { for it.Seek(b); it.Valid(); it.Next() { count++ } } if count > 0 { count-- } return nil }) return }) return } func (ncdbh *NodeChatDBHandler) revertPreviousId() (err error) { err = ncdbh.db(func(d *badger.DB) (err error) { err = d.View(func(txn *badger.Txn) error { opt := badger.DefaultIteratorOptions opt.Reverse = true it := txn.NewIterator(opt) defer it.Close() it.Rewind() if it.Valid() { item := it.Item() if err = item.Value(func(val []byte) (err error) { var chatMessage *ChatMessage if err = json.Unmarshal(val, &chatMessage); err != nil { return err } ncdbh.PreviousId = chatMessage.ID return }); err != nil { return err } } else { ncdbh.PreviousId = 1 } return nil }) return }) return } func (ncdbh *NodeChatDBHandler) AddNewChatMessage(chatMessage *ChatMessage) (err error) { b := make([]byte, bufferSize) ncdbh.PreviousId++ binary.BigEndian.PutUint64(b, ncdbh.PreviousId) chatMessage.ID = ncdbh.PreviousId bs, err := json.Marshal(chatMessage) if err != nil { return } if err = ncdbh.db(func(d *badger.DB) (err error) { err = d.Update(func(txn *badger.Txn) (err error) { if updateErr := txn.Set(b, bs); updateErr != nil { return updateErr } return nil }) return }); err != nil { return } return } func (ncdbh *NodeChatDBHandler) DeleteChatMessage(key uint64) (err error) { if err = ncdbh.db(func(d *badger.DB) (err error) { err = d.Update(func(txn *badger.Txn) (err error) { b := make([]byte, bufferSize) binary.BigEndian.PutUint64(b, key) if err = txn.Delete(b); err != nil { return } ncdbh.ItemCount-- return }) return }); err != nil { return } return } func (ncdbh *NodeChatDBHandler) DeleteChatFile(filename string, key uint64) (err error) { if err = ncdbh.DeleteChatMessage(key); err != nil { return } err = os.Remove(filepath.Join(dataPath, "data", "chats", ncdbh.ChatID, "__files__", filename)) return } func (ncdbh *NodeChatDBHandler) ListChatMessages(lastIndex int, limit int) (chatMessages []*ChatMessage, l int, done bool, err error) { err = ncdbh.db(func(d *badger.DB) (err error) { err = d.View(func(txn *badger.Txn) (err error) { opt := badger.DefaultIteratorOptions opt.Reverse = true it := txn.NewIterator(opt) defer it.Close() b := make([]byte, bufferSize) if lastIndex <= 0 { binary.BigEndian.PutUint64(b, uint64(ncdbh.PreviousId)) } else { binary.BigEndian.PutUint64(b, uint64(lastIndex)) } x := 0 defer func() { if x < limit { done = true } }() chatMessages = make([]*ChatMessage, 0) for it.Seek(b); it.Valid(); it.Next() { if x >= limit { break } item := it.Item() if err = item.Value(func(val []byte) (err error) { var chatMessage *ChatMessage if err = json.Unmarshal(val, &chatMessage); err != nil { return err } l = int(chatMessage.ID) chatMessages = append(chatMessages, chatMessage) return }); err != nil { return err } x++ } return }) return }) return } func (ncdbh *NodeChatDBHandler) ListChatFiles(lastIndex int, limit int) (chatMessages []*ChatFile, l int, err error) { err = ncdbh.db(func(d *badger.DB) (err error) { err = d.View(func(txn *badger.Txn) (err error) { opt := badger.DefaultIteratorOptions opt.Reverse = true it := txn.NewIterator(opt) defer it.Close() b := make([]byte, bufferSize) var li int if lastIndex <= 0 { binary.BigEndian.PutUint64(b, uint64(ncdbh.PreviousId)) li = int(ncdbh.PreviousId) } else { binary.BigEndian.PutUint64(b, uint64(lastIndex)) li = lastIndex } x := 0 y := 0 defer func() { if li > x { l = li - x } else if li == 0 { if ncdbh.PreviousId > uint64(x) { l = int(ncdbh.PreviousId) - x } else { l = 0 } } else { l = 0 } }() defer it.Close() chatMessages = make([]*ChatFile, 0) for it.Seek(b); it.Valid(); it.Next() { if y >= limit || x >= int(ncdbh.PreviousId) { break } item := it.Item() if err = item.Value(func(val []byte) (err error) { var chatMessage *ChatMessage if err = json.Unmarshal(val, &chatMessage); err != nil { return err } if chatMessage.File != nil { chatMessage.File.ID = chatMessage.ID chatMessages = append(chatMessages, chatMessage.File) l = lastIndex - x y++ } return }); err != nil { return err } x++ } return }) return }) return } func (ncdbh *NodeChatDBHandler) GetChatMessage(index uint64) (chatMessage *ChatMessage, err error) { err = ncdbh.db(func(d *badger.DB) (err error) { err = d.View(func(txn *badger.Txn) (err error) { b := make([]byte, bufferSize) binary.BigEndian.PutUint64(b, uint64(index)) item, err := txn.Get(b) if err != nil { return } err = item.Value(func(val []byte) error { return json.Unmarshal(val, &chatMessage) }) return }) return }) return } func (ncdbh *NodeChatDBHandler) ModifyChatMessage(key uint64, newContent string) (err error) { b := make([]byte, bufferSize) binary.BigEndian.PutUint64(b, key) chatMessage, err := ncdbh.GetChatMessage(key) if err != nil { return } chatMessage.Content = newContent bs, err := json.Marshal(chatMessage) if err != nil { return } if err = ncdbh.db(func(d *badger.DB) (err error) { err = d.Update(func(txn *badger.Txn) (err error) { if updateErr := txn.Set(b, bs); updateErr != nil { return updateErr } return nil }) return }); err != nil { return } return }