dispatch/storage/boltdb/boltdb.go

package boltdb

import (
	"bytes"
	"encoding/binary"
	"strconv"

	bolt "go.etcd.io/bbolt"

	"github.com/khlieng/dispatch/pkg/session"
	"github.com/khlieng/dispatch/storage"
)

var (
	bucketUsers    = []byte("Users")
	bucketServers  = []byte("Servers")
	bucketChannels = []byte("Channels")
	bucketMessages = []byte("Messages")
	bucketSessions = []byte("Sessions")
)

// BoltStore implements storage.Store, storage.MessageStore and storage.SessionStore
type BoltStore struct {
	db *bolt.DB
}

func New(path string) (*BoltStore, error) {
	db, err := bolt.Open(path, 0600, nil)
	if err != nil {
		return nil, err
	}

	db.Update(func(tx *bolt.Tx) error {
		tx.CreateBucketIfNotExists(bucketUsers)
		tx.CreateBucketIfNotExists(bucketServers)
		tx.CreateBucketIfNotExists(bucketChannels)
		tx.CreateBucketIfNotExists(bucketMessages)
		tx.CreateBucketIfNotExists(bucketSessions)
		return nil
	})

	return &BoltStore{
		db,
	}, nil
}

func (s *BoltStore) Close() {
	s.db.Close()
}

func (s *BoltStore) GetUsers() ([]*storage.User, error) {
	var users []*storage.User

	s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketUsers)

		return b.ForEach(func(k, v []byte) error {
			user := storage.User{
				IDBytes: make([]byte, 8),
			}
			user.Unmarshal(v)
			copy(user.IDBytes, k)

			users = append(users, &user)

			return nil
		})
	})

	return users, nil
}

func (s *BoltStore) SaveUser(user *storage.User) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketUsers)

		if user.ID == 0 {
			user.ID, _ = b.NextSequence()
			user.IDBytes = idToBytes(user.ID)
		}
		user.Username = strconv.FormatUint(user.ID, 10)

		data, err := user.Marshal(nil)
		if err != nil {
			return err
		}

		return b.Put(user.IDBytes, data)
	})
}

func (s *BoltStore) DeleteUser(user *storage.User) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketServers)
		c := b.Cursor()

		for k, _ := c.Seek(user.IDBytes); bytes.HasPrefix(k, user.IDBytes); k, _ = c.Next() {
			b.Delete(k)
		}

		b = tx.Bucket(bucketChannels)
		c = b.Cursor()

		for k, _ := c.Seek(user.IDBytes); bytes.HasPrefix(k, user.IDBytes); k, _ = c.Next() {
			b.Delete(k)
		}

		return tx.Bucket(bucketUsers).Delete(user.IDBytes)
	})
}

func (s *BoltStore) GetServer(user *storage.User, address string) (*storage.Server, error) {
	var server *storage.Server

	err := s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketServers)
		id := serverID(user, address)

		v := b.Get(id)
		if v == nil {
			return storage.ErrNotFound
		} else {
			server = &storage.Server{}
			server.Unmarshal(v)
			return nil
		}
	})

	return server, err
}

func (s *BoltStore) GetServers(user *storage.User) ([]*storage.Server, error) {
	var servers []*storage.Server

	s.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket(bucketServers).Cursor()

		for k, v := c.Seek(user.IDBytes); bytes.HasPrefix(k, user.IDBytes); k, v = c.Next() {
			server := storage.Server{}
			server.Unmarshal(v)
			servers = append(servers, &server)
		}

		return nil
	})

	return servers, nil
}

func (s *BoltStore) SaveServer(user *storage.User, server *storage.Server) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketServers)
		data, _ := server.Marshal(nil)

		return b.Put(serverID(user, server.Host), data)
	})
}

func (s *BoltStore) RemoveServer(user *storage.User, address string) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		serverID := serverID(user, address)
		tx.Bucket(bucketServers).Delete(serverID)

		b := tx.Bucket(bucketChannels)
		c := b.Cursor()

		for k, _ := c.Seek(serverID); bytes.HasPrefix(k, serverID); k, _ = c.Next() {
			b.Delete(k)
		}

		return nil
	})
}

func (s *BoltStore) SetNick(user *storage.User, nick, address string) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketServers)
		id := serverID(user, address)

		server := storage.Server{}
		v := b.Get(id)
		if v != nil {
			server.Unmarshal(v)
			server.Nick = nick

			data, _ := server.Marshal(nil)
			return b.Put(id, data)
		}

		return nil
	})
}

func (s *BoltStore) SetServerName(user *storage.User, name, address string) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketServers)
		id := serverID(user, address)

		server := storage.Server{}
		v := b.Get(id)
		if v != nil {
			server.Unmarshal(v)
			server.Name = name

			data, _ := server.Marshal(nil)
			return b.Put(id, data)
		}

		return nil
	})
}

func (s *BoltStore) GetChannels(user *storage.User) ([]*storage.Channel, error) {
	var channels []*storage.Channel

	s.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket(bucketChannels).Cursor()

		for k, v := c.Seek(user.IDBytes); bytes.HasPrefix(k, user.IDBytes); k, v = c.Next() {
			channel := storage.Channel{}
			channel.Unmarshal(v)
			channels = append(channels, &channel)
		}

		return nil
	})

	return channels, nil
}

func (s *BoltStore) AddChannel(user *storage.User, channel *storage.Channel) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketChannels)
		data, _ := channel.Marshal(nil)

		return b.Put(channelID(user, channel.Server, channel.Name), data)
	})
}

func (s *BoltStore) RemoveChannel(user *storage.User, server, channel string) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketChannels)
		id := channelID(user, server, channel)

		return b.Delete(id)
	})
}

func (s *BoltStore) LogMessage(message *storage.Message) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b, err := tx.Bucket(bucketMessages).CreateBucketIfNotExists([]byte(message.Server + ":" + message.To))
		if err != nil {
			return err
		}

		data, err := message.Marshal(nil)
		if err != nil {
			return err
		}

		return b.Put([]byte(message.ID), data)
	})
}

func (s *BoltStore) GetMessages(server, channel string, count int, fromID string) ([]storage.Message, bool, error) {
	messages := make([]storage.Message, count)
	hasMore := false

	s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketMessages).Bucket([]byte(server + ":" + channel))
		if b == nil {
			return nil
		}

		c := b.Cursor()

		if fromID != "" {
			c.Seek([]byte(fromID))

			for k, v := c.Prev(); count > 0 && k != nil; k, v = c.Prev() {
				count--
				messages[count].Unmarshal(v)
			}
		} else {
			for k, v := c.Last(); count > 0 && k != nil; k, v = c.Prev() {
				count--
				messages[count].Unmarshal(v)
			}
		}

		c.Next()
		k, _ := c.Prev()
		hasMore = k != nil

		return nil
	})

	if count == 0 {
		return messages, hasMore, nil
	} else if count < len(messages) {
		return messages[count:], hasMore, nil
	}

	return nil, false, nil
}

func (s *BoltStore) GetMessagesByID(server, channel string, ids []string) ([]storage.Message, error) {
	messages := make([]storage.Message, len(ids))

	err := s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketMessages).Bucket([]byte(server + ":" + channel))

		for i, id := range ids {
			messages[i].Unmarshal(b.Get([]byte(id)))
		}
		return nil
	})
	return messages, err
}

func (s *BoltStore) GetSessions() ([]*session.Session, error) {
	var sessions []*session.Session

	err := s.db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketSessions)

		return b.ForEach(func(_ []byte, v []byte) error {
			session := session.Session{}
			_, err := session.Unmarshal(v)
			sessions = append(sessions, &session)
			return err
		})
	})

	if err != nil {
		return nil, err
	}

	return sessions, nil
}

func (s *BoltStore) SaveSession(session *session.Session) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		b := tx.Bucket(bucketSessions)

		data, err := session.Marshal(nil)
		if err != nil {
			return err
		}

		return b.Put([]byte(session.Key()), data)
	})
}

func (s *BoltStore) DeleteSession(key string) error {
	return s.db.Batch(func(tx *bolt.Tx) error {
		return tx.Bucket(bucketSessions).Delete([]byte(key))
	})
}

func serverID(user *storage.User, address string) []byte {
	id := make([]byte, 8+len(address))
	copy(id, user.IDBytes)
	copy(id[8:], address)
	return id
}

func channelID(user *storage.User, server, channel string) []byte {
	id := make([]byte, 8+len(server)+1+len(channel))
	copy(id, user.IDBytes)
	copy(id[8:], server)
	copy(id[8+len(server)+1:], channel)
	return id
}

func idToBytes(i uint64) []byte {
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, i)
	return b
}

func idFromBytes(b []byte) uint64 {
	return binary.BigEndian.Uint64(b)
}