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()
}
