1
0
Fork 0
mirror of https://github.com/diamondburned/arikawa.git synced 2024-11-09 16:35:12 +00:00
arikawa/state/store_default.go

686 lines
12 KiB
Go
Raw Normal View History

2020-01-18 07:07:52 +00:00
package state
import (
"sort"
"sync"
"github.com/diamondburned/arikawa/discord"
)
// TODO: make an ExpiryStore
type DefaultStore struct {
*DefaultStoreOptions
self discord.User
// includes normal and private
2020-01-18 21:04:12 +00:00
privates map[discord.Snowflake]*discord.Channel // channelID:channel
guilds map[discord.Snowflake]*discord.Guild // guildID:guild
channels map[discord.Snowflake][]discord.Channel // guildID:channels
members map[discord.Snowflake][]discord.Member // guildID:members
presences map[discord.Snowflake][]discord.Presence // guildID:presences
messages map[discord.Snowflake][]discord.Message // channelID:messages
2020-01-18 07:07:52 +00:00
mut sync.Mutex
}
type DefaultStoreOptions struct {
MaxMessages uint // default 50
}
var _ Store = (*DefaultStore)(nil)
func NewDefaultStore(opts *DefaultStoreOptions) *DefaultStore {
if opts == nil {
opts = &DefaultStoreOptions{
MaxMessages: 50,
}
}
return &DefaultStore{
DefaultStoreOptions: opts,
privates: map[discord.Snowflake]*discord.Channel{},
guilds: map[discord.Snowflake]*discord.Guild{},
2020-01-18 21:04:12 +00:00
channels: map[discord.Snowflake][]discord.Channel{},
2020-01-18 21:40:09 +00:00
members: map[discord.Snowflake][]discord.Member{},
2020-01-18 21:04:12 +00:00
presences: map[discord.Snowflake][]discord.Presence{},
messages: map[discord.Snowflake][]discord.Message{},
2020-01-18 07:07:52 +00:00
}
}
func (s *DefaultStore) Reset() error {
s.mut.Lock()
defer s.mut.Unlock()
s.self = discord.User{}
s.privates = map[discord.Snowflake]*discord.Channel{}
s.guilds = map[discord.Snowflake]*discord.Guild{}
s.messages = map[discord.Snowflake][]discord.Message{}
return nil
}
////
func (s *DefaultStore) Self() (*discord.User, error) {
s.mut.Lock()
defer s.mut.Unlock()
if !s.self.ID.Valid() {
return nil, ErrStoreNotFound
}
return &s.self, nil
}
2020-01-18 07:10:42 +00:00
func (s *DefaultStore) SelfSet(me *discord.User) error {
2020-01-18 07:07:52 +00:00
s.mut.Lock()
s.self = *me
s.mut.Unlock()
return nil
}
////
func (s *DefaultStore) Channel(id discord.Snowflake) (*discord.Channel, error) {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
for _, chs := range s.channels {
for _, ch := range chs {
2020-01-18 07:07:52 +00:00
if ch.ID == id {
return &ch, nil
}
}
}
return nil, ErrStoreNotFound
}
func (s *DefaultStore) Channels(
guildID discord.Snowflake) ([]discord.Channel, error) {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
chs, ok := s.channels[guildID]
2020-01-18 07:07:52 +00:00
if !ok {
return nil, ErrStoreNotFound
}
return append([]discord.Channel{}, chs...), nil
2020-01-18 07:07:52 +00:00
}
func (s *DefaultStore) PrivateChannels() ([]discord.Channel, error) {
s.mut.Lock()
var chs = make([]discord.Channel, 0, len(s.privates))
for _, ch := range s.privates {
chs = append(chs, *ch)
}
s.mut.Unlock()
sort.Slice(chs, func(i, j int) bool {
// Latest first
return chs[i].LastMessageID > chs[j].LastMessageID
})
return chs, nil
}
func (s *DefaultStore) ChannelSet(channel *discord.Channel) error {
s.mut.Lock()
defer s.mut.Unlock()
switch channel.Type {
case discord.DirectMessage, discord.GroupDM:
s.privates[channel.ID] = channel
default:
2020-01-18 21:40:09 +00:00
chs := s.channels[channel.GuildID]
2020-01-18 07:07:52 +00:00
2020-01-18 21:04:12 +00:00
for i, ch := range chs {
2020-01-18 07:07:52 +00:00
if ch.ID == channel.ID {
// Also from discordgo.
if channel.Permissions == nil {
channel.Permissions = ch.Permissions
}
2020-01-18 07:07:52 +00:00
// Found, just edit
2020-01-18 21:04:12 +00:00
chs[i] = *channel
2020-01-18 07:07:52 +00:00
return nil
}
}
2020-01-18 21:04:12 +00:00
chs = append(chs, *channel)
s.channels[channel.GuildID] = chs
2020-01-18 07:07:52 +00:00
}
return nil
}
func (s *DefaultStore) ChannelRemove(channel *discord.Channel) error {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
chs, ok := s.channels[channel.GuildID]
2020-01-18 07:07:52 +00:00
if !ok {
return ErrStoreNotFound
}
2020-01-18 21:04:12 +00:00
for i, ch := range chs {
2020-01-18 07:07:52 +00:00
if ch.ID == channel.ID {
2020-01-18 21:04:12 +00:00
chs = append(chs[:i], chs[i+1:]...)
s.channels[channel.GuildID] = chs
2020-01-18 07:07:52 +00:00
return nil
}
}
return ErrStoreNotFound
}
////
func (s *DefaultStore) Emoji(
guildID, emojiID discord.Snowflake) (*discord.Emoji, error) {
s.mut.Lock()
defer s.mut.Unlock()
gd, ok := s.guilds[guildID]
if !ok {
return nil, ErrStoreNotFound
}
for _, emoji := range gd.Emojis {
if emoji.ID == emojiID {
return &emoji, nil
}
}
return nil, ErrStoreNotFound
}
func (s *DefaultStore) Emojis(
guildID discord.Snowflake) ([]discord.Emoji, error) {
s.mut.Lock()
defer s.mut.Unlock()
gd, ok := s.guilds[guildID]
if !ok {
return nil, ErrStoreNotFound
}
return append([]discord.Emoji{}, gd.Emojis...), nil
2020-01-18 07:07:52 +00:00
}
func (s *DefaultStore) EmojiSet(
guildID discord.Snowflake, emojis []discord.Emoji) error {
s.mut.Lock()
defer s.mut.Unlock()
gd, ok := s.guilds[guildID]
if !ok {
return ErrStoreNotFound
}
filtered := emojis[:0]
Main:
for _, enew := range emojis {
// Try and see if this emoji is already in the slice
for i, emoji := range gd.Emojis {
if emoji.ID == enew.ID {
// If it is, we simply replace it
gd.Emojis[i] = enew
continue Main
}
}
// If not, we add it to the slice that's to be appended.
filtered = append(filtered, enew)
}
// Append the new emojis
gd.Emojis = append(gd.Emojis, filtered...)
return nil
}
////
func (s *DefaultStore) Guild(id discord.Snowflake) (*discord.Guild, error) {
s.mut.Lock()
defer s.mut.Unlock()
ch, ok := s.guilds[id]
if !ok {
return nil, ErrStoreNotFound
}
return ch, nil
}
func (s *DefaultStore) Guilds() ([]discord.Guild, error) {
s.mut.Lock()
if len(s.guilds) == 0 {
s.mut.Unlock()
return nil, ErrStoreNotFound
}
2020-01-18 07:07:52 +00:00
var gs = make([]discord.Guild, 0, len(s.guilds))
for _, g := range s.guilds {
gs = append(gs, *g)
}
s.mut.Unlock()
sort.Slice(gs, func(i, j int) bool {
return gs[i].ID > gs[j].ID
})
return gs, nil
}
func (s *DefaultStore) GuildSet(guild *discord.Guild) error {
2020-01-18 07:07:52 +00:00
s.mut.Lock()
defer s.mut.Unlock()
if g, ok := s.guilds[guild.ID]; ok {
// preserve state stuff
if guild.Roles == nil {
guild.Roles = g.Roles
}
if guild.Emojis == nil {
guild.Emojis = g.Emojis
}
}
2020-01-18 21:04:12 +00:00
s.guilds[guild.ID] = guild
2020-01-18 07:07:52 +00:00
return nil
}
2020-01-18 21:04:12 +00:00
func (s *DefaultStore) GuildRemove(id discord.Snowflake) error {
2020-01-18 07:07:52 +00:00
s.mut.Lock()
2020-01-18 21:04:12 +00:00
delete(s.guilds, id)
2020-01-18 07:07:52 +00:00
s.mut.Unlock()
return nil
}
////
func (s *DefaultStore) Member(
guildID, userID discord.Snowflake) (*discord.Member, error) {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
ms, ok := s.members[guildID]
2020-01-18 07:07:52 +00:00
if !ok {
return nil, ErrStoreNotFound
}
2020-01-18 21:04:12 +00:00
for _, m := range ms {
if m.User.ID == userID {
return &m, nil
2020-01-18 07:07:52 +00:00
}
}
return nil, ErrStoreNotFound
}
func (s *DefaultStore) Members(
guildID discord.Snowflake) ([]discord.Member, error) {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
ms, ok := s.members[guildID]
2020-01-18 07:07:52 +00:00
if !ok {
return nil, ErrStoreNotFound
}
return append([]discord.Member{}, ms...), nil
2020-01-18 07:07:52 +00:00
}
func (s *DefaultStore) MemberSet(
guildID discord.Snowflake, member *discord.Member) error {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:40:09 +00:00
ms := s.members[guildID]
2020-01-18 07:07:52 +00:00
// Try and see if this member is already in the slice
2020-01-18 21:04:12 +00:00
for i, m := range ms {
2020-01-18 07:07:52 +00:00
if m.User.ID == member.User.ID {
// If it is, we simply replace it
2020-01-18 21:04:12 +00:00
ms[i] = *member
s.members[guildID] = ms
2020-01-18 07:07:52 +00:00
return nil
}
}
// Append the new member
2020-01-18 21:04:12 +00:00
ms = append(ms, *member)
s.members[guildID] = ms
2020-01-18 07:07:52 +00:00
return nil
}
func (s *DefaultStore) MemberRemove(guildID, userID discord.Snowflake) error {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
ms, ok := s.members[guildID]
2020-01-18 07:07:52 +00:00
if !ok {
return ErrStoreNotFound
}
// Try and see if this member is already in the slice
2020-01-18 21:04:12 +00:00
for i, m := range ms {
2020-01-18 07:07:52 +00:00
if m.User.ID == userID {
2020-01-18 21:04:12 +00:00
ms = append(ms, ms[i+1:]...)
s.members[guildID] = ms
2020-01-18 07:07:52 +00:00
return nil
}
}
return ErrStoreNotFound
}
////
func (s *DefaultStore) Message(
channelID, messageID discord.Snowflake) (*discord.Message, error) {
s.mut.Lock()
defer s.mut.Unlock()
ms, ok := s.messages[channelID]
if !ok {
return nil, ErrStoreNotFound
}
for _, m := range ms {
if m.ID == messageID {
return &m, nil
}
}
return nil, ErrStoreNotFound
}
func (s *DefaultStore) Messages(
channelID discord.Snowflake) ([]discord.Message, error) {
s.mut.Lock()
defer s.mut.Unlock()
ms, ok := s.messages[channelID]
if !ok {
return nil, ErrStoreNotFound
}
return append([]discord.Message{}, ms...), nil
2020-01-18 07:07:52 +00:00
}
func (s *DefaultStore) MaxMessages() int {
return int(s.DefaultStoreOptions.MaxMessages)
}
2020-01-18 07:07:52 +00:00
func (s *DefaultStore) MessageSet(message *discord.Message) error {
s.mut.Lock()
defer s.mut.Unlock()
ms, ok := s.messages[message.ChannelID]
if !ok {
ms = make([]discord.Message, 0, s.MaxMessages()+1)
2020-01-18 07:07:52 +00:00
}
// Check if we already have the message.
for i, m := range ms {
if m.ID == message.ID {
// Thanks, Discord.
if message.Content != "" {
m.Content = message.Content
}
if message.EditedTimestamp != nil {
m.EditedTimestamp = message.EditedTimestamp
}
if message.Mentions != nil {
m.Mentions = message.Mentions
}
if message.Embeds != nil {
m.Embeds = message.Embeds
}
if message.Attachments != nil {
m.Attachments = message.Attachments
}
if message.Timestamp.Valid() {
m.Timestamp = message.Timestamp
}
if message.Author.ID.Valid() {
m.Author = message.Author
}
2020-01-18 07:07:52 +00:00
ms[i] = m
return nil
}
}
// Prepend the latest message at the end
if len(ms) > 0 {
var end = s.MaxMessages()
if len(ms) < end {
end = len(ms)
}
// Copy hack to prepend. This copies the 0th-(end-1)th entries to
// 1st-endth.
copy(ms[1:end], ms[0:end-1])
// Then, set the 0th entry.
ms[0] = *message
2020-01-18 07:07:52 +00:00
} else {
ms = append(ms, *message)
2020-01-18 07:07:52 +00:00
}
2020-01-18 21:40:09 +00:00
s.messages[message.ChannelID] = ms
2020-01-18 07:07:52 +00:00
return nil
}
2020-01-18 21:04:12 +00:00
func (s *DefaultStore) MessageRemove(
channelID, messageID discord.Snowflake) error {
2020-01-18 07:07:52 +00:00
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
ms, ok := s.messages[channelID]
2020-01-18 07:07:52 +00:00
if !ok {
return ErrStoreNotFound
}
for i, m := range ms {
2020-01-18 21:04:12 +00:00
if m.ID == messageID {
2020-01-18 07:07:52 +00:00
ms = append(ms[:i], ms[i+1:]...)
2020-01-18 21:04:12 +00:00
s.messages[channelID] = ms
2020-01-18 07:07:52 +00:00
return nil
}
}
return ErrStoreNotFound
}
////
func (s *DefaultStore) Presence(
guildID, userID discord.Snowflake) (*discord.Presence, error) {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
ps, ok := s.presences[guildID]
2020-01-18 07:07:52 +00:00
if !ok {
return nil, ErrStoreNotFound
}
2020-01-18 21:04:12 +00:00
for _, p := range ps {
2020-01-18 07:07:52 +00:00
if p.User.ID == userID {
return &p, nil
}
}
return nil, ErrStoreNotFound
}
func (s *DefaultStore) Presences(
guildID discord.Snowflake) ([]discord.Presence, error) {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
ps, ok := s.presences[guildID]
2020-01-18 07:07:52 +00:00
if !ok {
return nil, ErrStoreNotFound
}
return append([]discord.Presence{}, ps...), nil
2020-01-18 07:07:52 +00:00
}
func (s *DefaultStore) PresenceSet(
guildID discord.Snowflake, presence *discord.Presence) error {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:40:09 +00:00
ps := s.presences[guildID]
2020-01-18 07:07:52 +00:00
2020-01-18 21:04:12 +00:00
for i, p := range ps {
2020-01-18 07:07:52 +00:00
if p.User.ID == presence.User.ID {
2020-01-18 21:04:12 +00:00
ps[i] = *presence
s.presences[guildID] = ps
2020-01-18 07:07:52 +00:00
return nil
}
}
2020-01-18 21:04:12 +00:00
ps = append(ps, *presence)
s.presences[guildID] = ps
2020-01-18 07:07:52 +00:00
return nil
}
func (s *DefaultStore) PresenceRemove(guildID, userID discord.Snowflake) error {
s.mut.Lock()
defer s.mut.Unlock()
2020-01-18 21:04:12 +00:00
ps, ok := s.presences[guildID]
2020-01-18 07:07:52 +00:00
if !ok {
return ErrStoreNotFound
}
2020-01-18 21:04:12 +00:00
for i, p := range ps {
2020-01-18 07:07:52 +00:00
if p.User.ID == userID {
2020-01-18 21:04:12 +00:00
ps = append(ps[:i], ps[i+1:]...)
s.presences[guildID] = ps
2020-01-18 07:07:52 +00:00
return nil
}
}
return ErrStoreNotFound
}
////
func (s *DefaultStore) Role(
guildID, roleID discord.Snowflake) (*discord.Role, error) {
s.mut.Lock()
defer s.mut.Unlock()
gd, ok := s.guilds[guildID]
if !ok {
return nil, ErrStoreNotFound
}
for _, r := range gd.Roles {
if r.ID == roleID {
return &r, nil
}
}
return nil, ErrStoreNotFound
}
func (s *DefaultStore) Roles(
guildID discord.Snowflake) ([]discord.Role, error) {
s.mut.Lock()
defer s.mut.Unlock()
gd, ok := s.guilds[guildID]
if !ok {
return nil, ErrStoreNotFound
}
return append([]discord.Role{}, gd.Roles...), nil
2020-01-18 07:07:52 +00:00
}
func (s *DefaultStore) RoleSet(
guildID discord.Snowflake, role *discord.Role) error {
s.mut.Lock()
defer s.mut.Unlock()
gd, ok := s.guilds[guildID]
if !ok {
return ErrStoreNotFound
}
for i, r := range gd.Roles {
if r.ID == role.ID {
gd.Roles[i] = *role
return nil
}
}
gd.Roles = append(gd.Roles, *role)
return nil
}
func (s *DefaultStore) RoleRemove(guildID, roleID discord.Snowflake) error {
s.mut.Lock()
defer s.mut.Unlock()
gd, ok := s.guilds[guildID]
if !ok {
return ErrStoreNotFound
}
for i, r := range gd.Roles {
if r.ID == roleID {
gd.Roles = append(gd.Roles[:i], gd.Roles[i+1:]...)
return nil
}
}
return ErrStoreNotFound
}