324 lines
7.5 KiB
Go
324 lines
7.5 KiB
Go
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
|
|
}
|