db: Reorganize code into separate package

db/db.go

@@ -1,4 +1,4 @@
-package main
+package db
 
 import (
 	"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 starred(id STRING NOT NULL UNIQUE)")
 }
+
+func Close() {
+	db.Close()
+}

db/prestige.go

@@ -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
+}

db/shard.go

@@ -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
+}

db/starboard.go

@@ -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
+}

db/xp.go

@@ -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
+}

gacha.go

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

levels.go

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

main.go

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

messageCreate.go

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

prestige.go

@@ -1,68 +1,9 @@
 package main
 
 import (
-	"sync"
-
 	"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) {
 	member, err := s.GuildMember(guildID, discordID)
 	if err != nil {

shard.go

@@ -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
-}

starboard.go

@@ -7,6 +7,8 @@ import (
 	"time"
 
 	"github.com/bwmarrin/discordgo"
+
+	"PermissionGacha/db"
 )
 
 const starEmoji = "⭐"
@@ -32,7 +34,12 @@ func listenToStarboardReact(s *discordgo.Session, m *discordgo.MessageReactionAd
 	starboardMutex.Lock()
 	defer starboardMutex.Unlock()
 
-	if checkStarredBefore(s, m.MessageID) {
+	starredBefore, err := db.HasStarredBefore(m.MessageID)
+	if err != nil {
+		logError(s, err)
+		return
+	}
+	if starredBefore {
 		return
 	}
 
@@ -90,25 +97,13 @@ func listenToStarboardReact(s *discordgo.Session, m *discordgo.MessageReactionAd
 		return
 	}
 	err = s.MessageReactionAdd(msg.ChannelID, msg.ID, starEmoji)
-	noteStarred(s, msg.ID)
 	if err != nil {
 		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 {
-		logError(s, "noteStarred", "exec", err)
+		logError(s, err)
+		return
 	}
 }

xp.go

@@ -1,15 +1,20 @@
 package main
 
 import (
+	"fmt"
 	"math"
 	"sync"
 	"time"
 
 	"github.com/bwmarrin/discordgo"
+
+	"PermissionGacha/db"
 )
 
 const notInGeneralNerf = 10
 
+const generalChannelID = "645550570001924098"
+
 var lastMessage = make(map[string]time.Time)
 var incrementMutex = sync.Mutex{}
 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 *= 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 {
-		logError(dg, "incrementXP", "insert", discordID, err)
+		logError(s, err)
-		tx.Rollback()
 		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 {
-		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 {
@@ -96,3 +76,12 @@ func convertXPIntoLevel(xp int) (int, int) {
 	}
 	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
+}