133 lines
3.1 KiB
Go
133 lines
3.1 KiB
Go
package badger
|
|
|
|
import (
|
|
"bytes"
|
|
"errors"
|
|
"fmt"
|
|
"strconv"
|
|
|
|
badger "github.com/dgraph-io/badger/v4"
|
|
log "github.com/sirupsen/logrus"
|
|
)
|
|
|
|
// IdsDB represents a Badger database
|
|
type IdsDB struct {
|
|
db *badger.DB
|
|
}
|
|
|
|
// IdsDBOpen returns a new DB object
|
|
func IdsDBOpen(path string) IdsDB {
|
|
bdb, err := badger.Open(badger.DefaultOptions(path))
|
|
if err != nil {
|
|
log.Errorf("Failed to open ids database: %v, falling back to in-memory database", path)
|
|
bdb, err = badger.Open(badger.DefaultOptions("").WithInMemory(true))
|
|
if err != nil {
|
|
log.Fatalf("Couldn't initialize the ids database")
|
|
}
|
|
}
|
|
|
|
return IdsDB{
|
|
db: bdb,
|
|
}
|
|
}
|
|
|
|
// Set stores an id pair
|
|
func (db *IdsDB) Set(tgAccount, xmppAccount string, tgChatId, tgMsgId int64, xmppId string) error {
|
|
bPrefix := toKeyPrefix(tgAccount, xmppAccount)
|
|
bTgId := toTgByteString(tgChatId, tgMsgId)
|
|
bXmppId := toXmppByteString(xmppId)
|
|
bTgKey := toByteKey(bPrefix, bTgId, "tg")
|
|
bXmppKey := toByteKey(bPrefix, bXmppId, "xmpp")
|
|
|
|
return db.db.Update(func(txn *badger.Txn) error {
|
|
if err := txn.Set(bTgKey, bXmppId); err != nil {
|
|
return err
|
|
}
|
|
return txn.Set(bXmppKey, bTgId)
|
|
})
|
|
}
|
|
|
|
func (db *IdsDB) getByteValue(key []byte) ([]byte, error) {
|
|
var valCopy []byte
|
|
err := db.db.View(func(txn *badger.Txn) error {
|
|
item, err := txn.Get(key)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
valCopy, err = item.ValueCopy(nil)
|
|
return err
|
|
})
|
|
return valCopy, err
|
|
}
|
|
|
|
// GetByTgIds obtains an XMPP id by Telegram chat/message ids
|
|
func (db *IdsDB) GetByTgIds(tgAccount, xmppAccount string, tgChatId, tgMsgId int64) (string, error) {
|
|
val, err := db.getByteValue(toByteKey(
|
|
toKeyPrefix(tgAccount, xmppAccount),
|
|
toTgByteString(tgChatId, tgMsgId),
|
|
"tg",
|
|
))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return string(val), nil
|
|
}
|
|
|
|
// GetByXmppId obtains Telegram chat/message ids by an XMPP id
|
|
func (db *IdsDB) GetByXmppId(tgAccount, xmppAccount, xmppId string) (int64, int64, error) {
|
|
val, err := db.getByteValue(toByteKey(
|
|
toKeyPrefix(tgAccount, xmppAccount),
|
|
toXmppByteString(xmppId),
|
|
"xmpp",
|
|
))
|
|
if err != nil {
|
|
return 0, 0, err
|
|
}
|
|
return splitTgByteString(val)
|
|
}
|
|
|
|
func toKeyPrefix(tgAccount, xmppAccount string) []byte {
|
|
return []byte(fmt.Sprintf("%v/%v/", tgAccount, xmppAccount))
|
|
}
|
|
|
|
func toByteKey(prefix, suffix []byte, typ string) []byte {
|
|
key := make([]byte, 0, len(prefix) + len(suffix) + 6)
|
|
key = append(key, prefix...)
|
|
key = append(key, []byte(typ)...)
|
|
key = append(key, []byte("/")...)
|
|
key = append(key, suffix...)
|
|
return key
|
|
}
|
|
|
|
func toTgByteString(tgChatId, tgMsgId int64) []byte {
|
|
return []byte(fmt.Sprintf("%v/%v", tgChatId, tgMsgId))
|
|
}
|
|
|
|
func toXmppByteString(xmppId string) []byte {
|
|
return []byte(xmppId)
|
|
}
|
|
|
|
func splitTgByteString(val []byte) (int64, int64, error) {
|
|
parts := bytes.Split(val, []byte("/"))
|
|
if len(parts) != 2 {
|
|
return 0, 0, errors.New("Couldn't parse tg id pair")
|
|
}
|
|
tgChatId, err := strconv.ParseInt(string(parts[0]), 10, 64)
|
|
if err != nil {
|
|
return 0, 0, err
|
|
}
|
|
tgMsgId, err := strconv.ParseInt(string(parts[1]), 10, 64)
|
|
return tgChatId, tgMsgId, err
|
|
}
|
|
|
|
// Gc compacts the value log
|
|
func (db *IdsDB) Gc() {
|
|
db.db.RunValueLogGC(0.7)
|
|
}
|
|
|
|
// Close closes a DB
|
|
func (db *IdsDB) Close() {
|
|
db.db.Close()
|
|
}
|