db: Reorganize code into separate package

master 0.7
Luther Wen Xu 2019-11-20 20:59:48 +07:00
parent 48b6994dc9
commit 4887d55707
Signed by: chanbakjsd
GPG Key ID: B7D77E3E9D102B70
13 changed files with 217 additions and 161 deletions

@ -1,4 +1,4 @@
package main package db
import ( import (
"database/sql" "database/sql"
@ -20,3 +20,7 @@ func init() {
db.Exec("CREATE TABLE IF NOT EXISTS prestige(id STRING NOT NULL UNIQUE, multiplier INTEGER NOT NULL)") db.Exec("CREATE TABLE IF NOT EXISTS prestige(id STRING NOT NULL UNIQUE, multiplier INTEGER NOT NULL)")
db.Exec("CREATE TABLE IF NOT EXISTS starred(id STRING NOT NULL UNIQUE)") db.Exec("CREATE TABLE IF NOT EXISTS starred(id STRING NOT NULL UNIQUE)")
} }
func Close() {
db.Close()
}

@ -0,0 +1,60 @@
package db
import (
"fmt"
"sync"
)
var (
prestigeCache = make(map[string]int)
prestigeMutex sync.Mutex
)
func SetPrestigeMultiplier(discordID string, newMultiplier int) error {
prestigeMutex.Lock()
defer prestigeMutex.Unlock()
tx, err := db.Begin()
if err != nil {
return fmt.Errorf("db: SetPrestige: underlying SQL error on 'begin': %w", err)
}
_, err = tx.Exec("REPLACE INTO prestige(id, multiplier) VALUES(?,?)", discordID, newMultiplier)
if err != nil {
tx.Rollback()
return fmt.Errorf("db: SetPrestige: underlying SQL error on 'replace/prestige': %w", err)
}
_, err = tx.Exec("REPLACE INTO xp(id, xp) VALUES(?,0)", discordID)
if err != nil {
tx.Rollback()
return fmt.Errorf("db: SetPrestige: underlying SQL error on 'replace/xp': %w", err)
}
err = tx.Commit()
if err != nil {
return fmt.Errorf("db: SetPrestige: underlying SQL error on 'commit': %w", err)
}
prestigeCache[discordID] = newMultiplier
return nil
}
func GetPrestigeMultiplier(discordID string) (int, error) {
prestigeMutex.Lock()
defer prestigeMutex.Unlock()
if prestigeCache[discordID] > 0 {
return prestigeCache[discordID], nil
}
rows, err := db.Query("SELECT multiplier FROM prestige WHERE id=?", discordID)
if err != nil {
return 1, fmt.Errorf("db: GetPrestigeMultiplier: underlying SQL error on 'select' (%s): %w", discordID, err)
}
defer rows.Close()
if rows.Next() {
var multiplier int
rows.Scan(&multiplier)
prestigeCache[discordID] = multiplier
return multiplier, nil
}
prestigeCache[discordID] = 1
return 1, nil
}

@ -0,0 +1,25 @@
package db
import "fmt"
func RewardShard(discordID string, count int) error {
_, err := db.Exec("INSERT INTO shard(id, shard) VALUES(?, 1) ON CONFLICT(id) DO UPDATE SET shard = shard+1", discordID)
if err != nil {
return fmt.Errorf("db: RewardShard: underlying SQL error on 'insert' (%s, %d): %w", discordID, count, err)
}
return nil
}
func GetShard(discordID string) (int, error) {
rows, err := db.Query("SELECT shard FROM shard WHERE id=?", discordID)
if err != nil {
return 0, fmt.Errorf("db: GetShard: underlying SQL error on 'select': %w", err)
}
defer rows.Close()
if rows.Next() {
var shard int
rows.Scan(&shard)
return shard, nil
}
return 0, nil
}

@ -0,0 +1,20 @@
package db
import "fmt"
func HasStarredBefore(discordID string) (bool, error) {
rows, err := db.Query("SELECT id FROM starred WHERE id=?", discordID)
if err != nil {
return true, fmt.Errorf("db: HasStarredBefore: underlying SQL error on 'select' (%s): %w", discordID, err)
}
defer rows.Close()
return rows.Next(), nil
}
func SetStarred(discordID string) error {
_, err := db.Exec("INSERT INTO starred(id) VALUES(?)", discordID)
if err != nil {
return fmt.Errorf("db: SetStarred: underlying SQL error on 'insert' (%s): %w", discordID, err)
}
return nil
}

@ -0,0 +1,28 @@
package db
import "fmt"
func IncrementXP(discordID string, amount int) error {
_, err := db.Exec(
"INSERT INTO xp(id, xp) VALUES(?, ?) ON CONFLICT(id) DO UPDATE SET xp = xp+?",
discordID, amount, amount,
)
if err != nil {
return fmt.Errorf("db: IncrementXP: underlying SQL error on 'insert' (%s, %d): %w", discordID, amount, err)
}
return nil
}
func GetXP(discordID string) (int, error) {
rows, err := db.Query("SELECT xp FROM xp WHERE id=?", discordID)
if err != nil {
return 0, fmt.Errorf("db: GetXP: underlying SQL error on 'select' (%s): %w", discordID, err)
}
defer rows.Close()
if rows.Next() {
var xp int
rows.Scan(&xp)
return xp, nil
}
return 0, nil
}

@ -5,6 +5,8 @@ import (
"fmt" "fmt"
"github.com/bwmarrin/discordgo" "github.com/bwmarrin/discordgo"
"PermissionGacha/db"
) )
const ( const (
@ -53,7 +55,7 @@ func rollGacha(dg *discordgo.Session, member *discordgo.Member) {
} }
channel, err := dg.UserChannelCreate(member.User.ID) channel, err := dg.UserChannelCreate(member.User.ID)
if found { if found {
rewardShard(dg, member.User.ID) db.RewardShard(member.User.ID, 1)
if err == nil { if err == nil {
dg.ChannelMessageSend( dg.ChannelMessageSend(
channel.ID, channel.ID,

@ -39,7 +39,11 @@ func listenToLevelUpdate(dg *discordgo.Session) {
continue continue
} }
levelQueue[k] = false levelQueue[k] = false
level, _ := getLevelAndXP(dg, k) level, _, err := getLevelAndXP(k)
if err != nil {
logError(dg, err)
continue
}
if assignedLevels[k] != level { if assignedLevels[k] != level {
assignedLevels[k] = level assignedLevels[k] = level
go checkForGachaEligibility(dg, k, rolesFromLevel[level-1]) go checkForGachaEligibility(dg, k, rolesFromLevel[level-1])

@ -9,6 +9,8 @@ import (
"syscall" "syscall"
"github.com/bwmarrin/discordgo" "github.com/bwmarrin/discordgo"
"PermissionGacha/db"
) )
func main() { func main() {

@ -8,9 +8,10 @@ import (
"time" "time"
"github.com/bwmarrin/discordgo" "github.com/bwmarrin/discordgo"
"PermissionGacha/db"
) )
const generalChannelID = "645550570001924098"
const emojiCheck = "✅" const emojiCheck = "✅"
const prestigeRequirement = 90000 const prestigeRequirement = 90000
@ -20,7 +21,7 @@ func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
if m.Author.Bot { if m.Author.Bot {
return return
} }
go incrementXP(s, m.Author.ID, m.ChannelID == generalChannelID) go incrementXP(s, m)
if m.Content == "pgqlevel" { if m.Content == "pgqlevel" {
handleXPRequest(s, m) handleXPRequest(s, m)
@ -37,7 +38,11 @@ func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
} }
func handleXPRequest(s *discordgo.Session, m *discordgo.MessageCreate) { func handleXPRequest(s *discordgo.Session, m *discordgo.MessageCreate) {
level, rawXP := getLevelAndXP(s, m.Author.ID) level, rawXP, err := getLevelAndXP(m.Author.ID)
if err != nil {
logError(s, err)
return
}
randomSource := make([]byte, 1) randomSource := make([]byte, 1)
rand.Read(randomSource) rand.Read(randomSource)
xp := rawXP - (rawXP % 100) + ((int(randomSource[0]) + 128) % 100) xp := rawXP - (rawXP % 100) + ((int(randomSource[0]) + 128) % 100)
@ -55,7 +60,11 @@ func handleXPRequest(s *discordgo.Session, m *discordgo.MessageCreate) {
} }
func handleShardRequest(s *discordgo.Session, m *discordgo.MessageCreate) { func handleShardRequest(s *discordgo.Session, m *discordgo.MessageCreate) {
shard := getShard(s, m.Author.ID) shard, err := db.GetShard(m.Author.ID)
if err != nil {
logError(s, err)
return
}
s.ChannelMessageSend( s.ChannelMessageSend(
m.ChannelID, m.ChannelID,
fmt.Sprintf( fmt.Sprintf(
@ -66,8 +75,16 @@ func handleShardRequest(s *discordgo.Session, m *discordgo.MessageCreate) {
} }
func handlePrestigeRequest(s *discordgo.Session, m *discordgo.MessageCreate) { func handlePrestigeRequest(s *discordgo.Session, m *discordgo.MessageCreate) {
currentPrestigeLevel := getPrestige(s, m.Author.ID) currentPrestigeLevel, err := db.GetPrestigeMultiplier(m.Author.ID)
xp := getXP(s, m.Author.ID) if err != nil {
logError(s, err)
return
}
xp, err := db.GetXP(m.Author.ID)
if err != nil {
logError(s, err)
return
}
newPrestigeLevel := xp/prestigeRequirement + 1 newPrestigeLevel := xp/prestigeRequirement + 1
if xp < levelUpRequirementCache[3] && currentPrestigeLevel == 1 { if xp < levelUpRequirementCache[3] && currentPrestigeLevel == 1 {
@ -120,10 +137,18 @@ func checkPrestigeReact(s *discordgo.Session, m *discordgo.MessageReactionAdd) {
if m.Emoji.Name != emojiCheck { if m.Emoji.Name != emojiCheck {
return return
} }
xp := getXP(s, m.UserID) xp, err := db.GetXP(m.UserID)
if err != nil {
logError(s, err)
return
}
newPrestigeLevel := xp/prestigeRequirement + 1 newPrestigeLevel := xp/prestigeRequirement + 1
go cleanUpRole(s, m.UserID) go cleanUpRole(s, m.UserID)
setPrestige(s, m.UserID, newPrestigeLevel) err = db.SetPrestigeMultiplier(m.UserID, newPrestigeLevel)
if err != nil {
logError(s, err)
return
}
s.ChannelMessageSend(m.ChannelID, s.ChannelMessageSend(m.ChannelID,
fmt.Sprintf( fmt.Sprintf(
"**PRESTIGE > ** <@%s>, you have prestiged!", m.UserID, "**PRESTIGE > ** <@%s>, you have prestiged!", m.UserID,

@ -1,68 +1,9 @@
package main package main
import ( import (
"sync"
"github.com/bwmarrin/discordgo" "github.com/bwmarrin/discordgo"
) )
var prestigeCache = make(map[string]int)
var prestigeMutex = sync.Mutex{}
func setPrestige(dg *discordgo.Session, discordID string, newMultiplier int) {
prestigeMutex.Lock()
defer prestigeMutex.Unlock()
tx, err := db.Begin()
if err != nil {
logError(dg, "setPrestige", "begin", discordID, newMultiplier, err)
return
}
_, err = tx.Exec("REPLACE INTO prestige(id, multiplier) VALUES(?,?)", discordID, newMultiplier)
if err != nil {
logError(dg, "setPrestige", "insertMultiplier", discordID, newMultiplier, err)
tx.Rollback()
return
}
_, err = tx.Exec("REPLACE INTO xp(id, xp) VALUES(?,0)", discordID)
err = tx.Commit()
if err != nil {
logError(dg, "setPrestige", "insertXP", discordID, newMultiplier, err)
tx.Rollback()
return
}
if err != nil {
logError(dg, "setPrestige", "commit", discordID, newMultiplier, err)
return
}
prestigeCache[discordID] = newMultiplier
}
func getPrestige(dg *discordgo.Session, discordID string) int {
prestigeMutex.Lock()
defer prestigeMutex.Unlock()
if prestigeCache[discordID] > 0 {
return prestigeCache[discordID]
}
rows, err := db.Query("SELECT multiplier FROM prestige WHERE id=?", discordID)
if err != nil {
logError(dg, "getPrestige", "select", discordID, err)
return 0
}
defer rows.Close()
if rows.Next() {
var multiplier int
rows.Scan(&multiplier)
prestigeCache[discordID] = multiplier
return multiplier
}
prestigeCache[discordID] = 1
return 1
}
func cleanUpRole(s *discordgo.Session, discordID string) { func cleanUpRole(s *discordgo.Session, discordID string) {
member, err := s.GuildMember(guildID, discordID) member, err := s.GuildMember(guildID, discordID)
if err != nil { if err != nil {

@ -1,39 +0,0 @@
package main
import (
"github.com/bwmarrin/discordgo"
)
func rewardShard(dg *discordgo.Session, discordID string) {
tx, err := db.Begin()
if err != nil {
logError(dg, "rewardShard", "begin", discordID, err)
return
}
_, err = tx.Exec("INSERT INTO shard(id, shard) VALUES(?, 1) ON CONFLICT(id) DO UPDATE SET shard = shard+1;", discordID)
if err != nil {
logError(dg, "rewardShard", "insert", discordID, err)
tx.Rollback()
return
}
err = tx.Commit()
if err != nil {
logError(dg, "rewardShard", "commit", discordID, err)
return
}
}
func getShard(dg *discordgo.Session, discordID string) int {
rows, err := db.Query("SELECT shard FROM shard WHERE id=?", discordID)
if err != nil {
logError(dg, "getShard", "select", discordID, err)
return 0
}
defer rows.Close()
if rows.Next() {
var shard int
rows.Scan(&shard)
return shard
}
return 0
}

@ -7,6 +7,8 @@ import (
"time" "time"
"github.com/bwmarrin/discordgo" "github.com/bwmarrin/discordgo"
"PermissionGacha/db"
) )
const starEmoji = "⭐" const starEmoji = "⭐"
@ -32,7 +34,12 @@ func listenToStarboardReact(s *discordgo.Session, m *discordgo.MessageReactionAd
starboardMutex.Lock() starboardMutex.Lock()
defer starboardMutex.Unlock() defer starboardMutex.Unlock()
if checkStarredBefore(s, m.MessageID) { starredBefore, err := db.HasStarredBefore(m.MessageID)
if err != nil {
logError(s, err)
return
}
if starredBefore {
return return
} }
@ -90,25 +97,13 @@ func listenToStarboardReact(s *discordgo.Session, m *discordgo.MessageReactionAd
return return
} }
err = s.MessageReactionAdd(msg.ChannelID, msg.ID, starEmoji) err = s.MessageReactionAdd(msg.ChannelID, msg.ID, starEmoji)
noteStarred(s, msg.ID)
if err != nil { if err != nil {
logError(s, "listenToStarboardReact", "messageReactionAdd", err) logError(s, "listenToStarboardReact", "messageReactionAdd", err)
return
} }
} err = db.SetStarred(msg.ID)
func checkStarredBefore(s *discordgo.Session, discordID string) bool {
rows, err := db.Query("SELECT id FROM starred WHERE id=?", discordID)
if err != nil {
logError(s, "checkStarredBefore", "query", err)
return true
}
defer rows.Close()
return rows.Next()
}
func noteStarred(s *discordgo.Session, discordID string) {
_, err := db.Exec("INSERT INTO starred(id) VALUES(?)", discordID)
if err != nil { if err != nil {
logError(s, "noteStarred", "exec", err) logError(s, err)
return
} }
} }

59
xp.go

@ -1,15 +1,20 @@
package main package main
import ( import (
"fmt"
"math" "math"
"sync" "sync"
"time" "time"
"github.com/bwmarrin/discordgo" "github.com/bwmarrin/discordgo"
"PermissionGacha/db"
) )
const notInGeneralNerf = 10 const notInGeneralNerf = 10
const generalChannelID = "645550570001924098"
var lastMessage = make(map[string]time.Time) var lastMessage = make(map[string]time.Time)
var incrementMutex = sync.Mutex{} var incrementMutex = sync.Mutex{}
var levelUpRequirementCache [30]int var levelUpRequirementCache [30]int
@ -23,49 +28,24 @@ func init() {
} }
} }
func incrementXP(dg *discordgo.Session, discordID string, inGeneral bool) { func incrementXP(s *discordgo.Session, m *discordgo.MessageCreate) {
amountToIncrement := calculateIncrement(discordID) amountToIncrement := calculateIncrement(m.Author.ID)
if !inGeneral { if m.ChannelID != generalChannelID {
amountToIncrement /= notInGeneralNerf amountToIncrement /= notInGeneralNerf
} }
amountToIncrement *= getPrestige(dg, discordID) multiplier, err := db.GetPrestigeMultiplier(m.Author.ID)
tx, err := db.Begin()
if err != nil {
logError(dg, "incrementXP", "begin", discordID, err)
return
}
_, err = tx.Exec("INSERT INTO xp(id, xp) VALUES(?, ?) ON CONFLICT(id) DO UPDATE SET xp = xp+?;", discordID, amountToIncrement, amountToIncrement)
if err != nil { if err != nil {
logError(dg, "incrementXP", "insert", discordID, err) logError(s, err)
tx.Rollback()
return return
} }
err = tx.Commit() amountToIncrement *= multiplier
if err != nil {
logError(dg, "incrementXP", "commit", discordID, err)
return
}
addToLevelQueue(discordID)
}
func getXP(dg *discordgo.Session, discordID string) int { err = db.IncrementXP(m.Author.ID, amountToIncrement)
rows, err := db.Query("SELECT xp FROM xp WHERE id=?", discordID)
if err != nil { if err != nil {
logError(dg, "getLevelAndXP", "select", discordID, err) logError(s, err)
return 0 return
}
defer rows.Close()
if rows.Next() {
var xp int
rows.Scan(&xp)
return xp
}
return 0
} }
addToLevelQueue(m.Author.ID)
func getLevelAndXP(dg *discordgo.Session, discordID string) (int, int) {
return convertXPIntoLevel(getXP(dg, discordID))
} }
func calculateIncrement(discordID string) int { func calculateIncrement(discordID string) int {
@ -96,3 +76,12 @@ func convertXPIntoLevel(xp int) (int, int) {
} }
return 30, xp - levelUpRequirementCache[29] return 30, xp - levelUpRequirementCache[29]
} }
func getLevelAndXP(discordID string) (int, int, error) {
rawXP, err := db.GetXP(discordID)
if err != nil {
return 1, 0, fmt.Errorf("xp: getLevelAndXP: error requesting XP: %w", err)
}
level, xp := convertXPIntoLevel(rawXP)
return level, xp, nil
}