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 SquadChatDBHandler struct {
|
|
SquadID string
|
|
PreviousId uint64
|
|
ItemCount uint64
|
|
db func(func(*badger.DB) (err error)) (err error)
|
|
lock *sync.RWMutex
|
|
}
|
|
|
|
func NewSquadChatDBHandler(squadId string) (squadChatDBHandler *SquadChatDBHandler, err error) {
|
|
squadChatDBHandler = &SquadChatDBHandler{
|
|
db: func(f func(*badger.DB) (err error)) (err error) {
|
|
db, err := badger.Open(badger.DefaultOptions(filepath.Join(dataPath, "data", "squads", squadId, "chat")).WithLogger(dbLogger))
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer db.Close()
|
|
err = f(db)
|
|
return
|
|
},
|
|
|
|
SquadID: squadId,
|
|
lock: new(sync.RWMutex),
|
|
}
|
|
err = squadChatDBHandler.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
|
|
}
|
|
}
|
|
squadChatDBHandler.PreviousId = counter
|
|
squadChatDBHandler.ItemCount = itemCounter
|
|
return nil
|
|
})
|
|
return
|
|
})
|
|
return
|
|
}
|
|
|
|
func (zcdbh *SquadChatDBHandler) calculateNewChatCount(previousId uint64) (count uint, err error) {
|
|
err = zcdbh.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 (zcdbh *SquadChatDBHandler) revertPreviousId() (err error) {
|
|
err = zcdbh.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
|
|
}
|
|
zcdbh.PreviousId = chatMessage.ID
|
|
return
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
zcdbh.PreviousId = 1
|
|
}
|
|
return nil
|
|
})
|
|
return
|
|
})
|
|
return
|
|
}
|
|
|
|
func (zcdbh *SquadChatDBHandler) AddNewChatMessage(chatMessage *ChatMessage) (err error) {
|
|
b := make([]byte, bufferSize)
|
|
zcdbh.PreviousId++
|
|
binary.BigEndian.PutUint64(b, zcdbh.PreviousId)
|
|
chatMessage.ID = zcdbh.PreviousId
|
|
bs, err := json.Marshal(chatMessage)
|
|
if err != nil {
|
|
return
|
|
}
|
|
if err = zcdbh.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 (zcdbh *SquadChatDBHandler) DeleteChatMessage(key uint64) (err error) {
|
|
if err = zcdbh.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
|
|
}
|
|
zcdbh.ItemCount--
|
|
return
|
|
})
|
|
return
|
|
}); err != nil {
|
|
return
|
|
}
|
|
return
|
|
}
|
|
|
|
func (zcdbh *SquadChatDBHandler) DeleteChatFile(filename string, key uint64) (err error) {
|
|
if err = zcdbh.DeleteChatMessage(key); err != nil {
|
|
return
|
|
}
|
|
err = os.Remove(filepath.Join(dataPath, "data", "squads", zcdbh.SquadID, "chat", "__files__", filename))
|
|
return
|
|
}
|
|
|
|
func (zcdbh *SquadChatDBHandler) ListChatMessages(lastIndex int, limit int) (chatMessages []*ChatMessage, l int, done bool, err error) {
|
|
err = zcdbh.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(zcdbh.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 (zcdbh *SquadChatDBHandler) ListChatFiles(lastIndex int, limit int) (chatMessages []*ChatFile, l int, err error) {
|
|
err = zcdbh.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(zcdbh.PreviousId))
|
|
li = int(zcdbh.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 zcdbh.PreviousId > uint64(x) {
|
|
l = int(zcdbh.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(zcdbh.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 {
|
|
chatMessages = append(chatMessages, chatMessage.File)
|
|
l = lastIndex - x
|
|
y++
|
|
}
|
|
return
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
x++
|
|
}
|
|
return
|
|
})
|
|
return
|
|
})
|
|
return
|
|
}
|
|
|
|
func (zcdbh *SquadChatDBHandler) GetChatMessage(index uint64) (chatMessage *ChatMessage, err error) {
|
|
err = zcdbh.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 (zcdbh *SquadChatDBHandler) ModifyChatMessage(key uint64, newContent string) (err error) {
|
|
b := make([]byte, bufferSize)
|
|
binary.BigEndian.PutUint64(b, key)
|
|
chatMessage, err := zcdbh.GetChatMessage(key)
|
|
if err != nil {
|
|
return
|
|
}
|
|
chatMessage.Content = newContent
|
|
bs, err := json.Marshal(chatMessage)
|
|
if err != nil {
|
|
return
|
|
}
|
|
if err = zcdbh.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
|
|
}
|