From f69b3f13a3472c6d59558f67a8e17f19791d5773 Mon Sep 17 00:00:00 2001 From: Siwoo Jeon Date: Sun, 11 May 2025 13:48:53 +0900 Subject: [PATCH] chore: edit database struct --- commands/dataLength.go | 10 ++++----- commands/deleteLearnedData.go | 2 +- commands/learn.go | 2 +- commands/learnedDataList.go | 2 +- components/deleteLearnedData.go | 2 +- databases/Learn.go | 4 ---- databases/Text.go | 4 ---- databases/database.go | 28 ++++++++++++++++++++----- handler/messageCreate.go | 8 ++++---- main.go | 2 +- scripts/export.go | 36 +++++---------------------------- 11 files changed, 42 insertions(+), 58 deletions(-) diff --git a/commands/dataLength.go b/commands/dataLength.go index e14ffaa..6ee92eb 100644 --- a/commands/dataLength.go +++ b/commands/dataLength.go @@ -89,9 +89,9 @@ func dataLengthRun(s *discordgo.Session, m any) { } dataLengthWg.Add(5) - go getLength(text, databases.Texts, bson.D{{}}) - go getLength(muffin, databases.Texts, bson.D{{Key: "persona", Value: "muffin"}}) - go getLength(nsfw, databases.Texts, bson.D{ + go getLength(text, databases.Database.Texts, bson.D{{}}) + go getLength(muffin, databases.Database.Texts, bson.D{{Key: "persona", Value: "muffin"}}) + go getLength(nsfw, databases.Database.Texts, bson.D{ { Key: "persona", Value: bson.M{ @@ -99,8 +99,8 @@ func dataLengthRun(s *discordgo.Session, m any) { }, }, }) - go getLength(learn, databases.Learns, bson.D{{}}) - go getLength(userLearn, databases.Learns, bson.D{{Key: "user_id", Value: userId}}) + go getLength(learn, databases.Database.Learns, bson.D{{}}) + go getLength(userLearn, databases.Database.Learns, bson.D{{Key: "user_id", Value: userId}}) go func() { dataLengthWg.Wait() diff --git a/commands/deleteLearnedData.go b/commands/deleteLearnedData.go index 105c928..aaa31e0 100644 --- a/commands/deleteLearnedData.go +++ b/commands/deleteLearnedData.go @@ -74,7 +74,7 @@ func deleteLearnedDataRun(c *Command, s *discordgo.Session, m any, args *[]strin userId = m.Member.User.ID } - cur, err := databases.Learns.Find(context.TODO(), bson.M{"user_id": userId, "command": command}) + cur, err := databases.Database.Learns.Find(context.TODO(), bson.M{"user_id": userId, "command": command}) if err != nil { embed := &discordgo.MessageEmbed{ Title: "❌ 오류", diff --git a/commands/learn.go b/commands/learn.go index 12af361..be0fe30 100644 --- a/commands/learn.go +++ b/commands/learn.go @@ -173,7 +173,7 @@ func learnRun(c *Command, s *discordgo.Session, m any, args *[]string) { } } - _, err := databases.Learns.InsertOne(context.TODO(), databases.InsertLearn{ + _, err := databases.Database.Learns.InsertOne(context.TODO(), databases.InsertLearn{ Command: command, Result: result, UserId: userId, diff --git a/commands/learnedDataList.go b/commands/learnedDataList.go index be71433..dc7b486 100644 --- a/commands/learnedDataList.go +++ b/commands/learnedDataList.go @@ -115,7 +115,7 @@ func learnedDataListRun(s *discordgo.Session, m any, args *[]string) { } } - cur, err := databases.Learns.Find(context.TODO(), bson.D{{Key: "user_id", Value: userId}, filter}) + cur, err := databases.Database.Learns.Find(context.TODO(), bson.D{{Key: "user_id", Value: userId}, filter}) if err != nil { if err == mongo.ErrNoDocuments { embed := &discordgo.MessageEmbed{ diff --git a/components/deleteLearnedData.go b/components/deleteLearnedData.go index 4737d11..84311a4 100644 --- a/components/deleteLearnedData.go +++ b/components/deleteLearnedData.go @@ -69,7 +69,7 @@ var DeleteLearnedDataComponent *commands.Component = &commands.Component{ id, _ := bson.ObjectIDFromHex(strings.ReplaceAll(utils.ItemIdRegexp.ReplaceAllString(i.MessageComponentData().Values[0][len(utils.DeleteLearnedData):], ""), "&", "")) itemId := strings.ReplaceAll(utils.ItemIdRegexp.FindAllString(i.MessageComponentData().Values[0], 1)[0], "No.", "") - databases.Learns.DeleteOne(context.TODO(), bson.D{{Key: "_id", Value: id}}) + databases.Database.Learns.DeleteOne(context.TODO(), bson.D{{Key: "_id", Value: id}}) i.EditReply(&discordgo.WebhookEdit{ Embeds: &[]*discordgo.MessageEmbed{ diff --git a/databases/Learn.go b/databases/Learn.go index 81ce270..cf0c265 100644 --- a/databases/Learn.go +++ b/databases/Learn.go @@ -3,9 +3,7 @@ package databases import ( "time" - "git.wh64.net/muffin/goMuffin/configs" "go.mongodb.org/mongo-driver/v2/bson" - "go.mongodb.org/mongo-driver/v2/mongo" ) type InsertLearn struct { @@ -22,5 +20,3 @@ type Learn struct { UserId string `bson:"user_id" json:"user_id"` CreatedAt time.Time `bson:"created_at" json:"created_at"` } - -var Learns *mongo.Collection = Client.Database(configs.Config.DBName).Collection("learn") diff --git a/databases/Text.go b/databases/Text.go index fbff6b2..7bacbbe 100644 --- a/databases/Text.go +++ b/databases/Text.go @@ -3,9 +3,7 @@ package databases import ( "time" - "git.wh64.net/muffin/goMuffin/configs" "go.mongodb.org/mongo-driver/v2/bson" - "go.mongodb.org/mongo-driver/v2/mongo" ) type InsertText struct { @@ -20,5 +18,3 @@ type Text struct { Persona string `bson:"persona" json:"persona"` CreatedAt time.Time `bson:"created_at" json:"created_at"` } - -var Texts *mongo.Collection = Client.Database(configs.Config.DBName).Collection("text") diff --git a/databases/database.go b/databases/database.go index 81115f2..7928bdc 100644 --- a/databases/database.go +++ b/databases/database.go @@ -8,13 +8,31 @@ import ( "go.mongodb.org/mongo-driver/v2/mongo/options" ) -func connect() *mongo.Client { - client, err := mongo.Connect(options.Client().ApplyURI(configs.Config.DatabaseURL)) +type MuffinDatabase struct { + Client *mongo.Client + Learns *mongo.Collection + Texts *mongo.Collection +} + +var Database *MuffinDatabase + +func init() { + var err error + + Database, err = Connect() if err != nil { log.Fatalln(err) } - - return client } -var Client *mongo.Client = connect() +func Connect() (*MuffinDatabase, error) { + client, err := mongo.Connect(options.Client().ApplyURI(configs.Config.DatabaseURL)) + if err != nil { + return nil, err + } + return &MuffinDatabase{ + Client: client, + Learns: client.Database(configs.Config.DBName).Collection("learn"), + Texts: client.Database(configs.Config.DBName).Collection("text"), + }, nil +} diff --git a/handler/messageCreate.go b/handler/messageCreate.go index d9d98bb..6832252 100644 --- a/handler/messageCreate.go +++ b/handler/messageCreate.go @@ -60,7 +60,7 @@ func MessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) { command := commands.Discommand.Aliases[args[0]] if m.Author.ID == config.Train.UserID { - if _, err := databases.Texts.InsertOne(context.TODO(), databases.InsertText{ + if _, err := databases.Database.Texts.InsertOne(context.TODO(), databases.InsertText{ Text: content, Persona: "muffin", CreatedAt: time.Now(), @@ -83,7 +83,7 @@ func MessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) { if channel.NSFW { filter = bson.D{{}} - if _, err := databases.Texts.InsertOne(context.TODO(), databases.InsertText{ + if _, err := databases.Database.Texts.InsertOne(context.TODO(), databases.InsertText{ Text: content, Persona: fmt.Sprintf("user:%s", m.Author.Username), CreatedAt: time.Now(), @@ -96,7 +96,7 @@ func MessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) { } go func() { - cur, err := databases.Texts.Find(context.TODO(), filter) + cur, err := databases.Database.Texts.Find(context.TODO(), filter) if err != nil { log.Fatalln(err) } @@ -107,7 +107,7 @@ func MessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) { ch <- 1 }() go func() { - cur, err := databases.Learns.Find(context.TODO(), bson.D{{Key: "command", Value: content}}) + cur, err := databases.Database.Learns.Find(context.TODO(), bson.D{{Key: "command", Value: content}}) if err != nil { if err == mongo.ErrNilDocument { learnData = []databases.Learn{} diff --git a/main.go b/main.go index e53f1f5..0ce9876 100644 --- a/main.go +++ b/main.go @@ -108,7 +108,7 @@ func main() { go dg.ApplicationCommandCreate(dg.State.User.ID, "", cmd.ApplicationCommand) } - defer databases.Client.Disconnect(context.TODO()) + defer databases.Database.Client.Disconnect(context.TODO()) log.Println("[goMuffin] 봇이 실행되고 있어요. 버전:", configs.MUFFIN_VERSION) sc := make(chan os.Signal, 1) diff --git a/scripts/export.go b/scripts/export.go index c8269c6..24c26b0 100644 --- a/scripts/export.go +++ b/scripts/export.go @@ -10,14 +10,11 @@ import ( "sync" "time" - "git.wh64.net/muffin/goMuffin/configs" "git.wh64.net/muffin/goMuffin/databases" "git.wh64.net/muffin/goMuffin/utils" "github.com/devproje/commando" "github.com/devproje/commando/option" "go.mongodb.org/mongo-driver/v2/bson" - "go.mongodb.org/mongo-driver/v2/mongo" - "go.mongodb.org/mongo-driver/v2/mongo/options" ) var date time.Time = time.Now() @@ -130,10 +127,11 @@ func saveFileToJSONL(path, name string, data any) error { } func ExportData(n *commando.Node) error { + defer databases.Database.Client.Disconnect(context.TODO()) + var wg sync.WaitGroup ch := make(chan error, 3) - databases.Client.Disconnect(context.TODO()) // databases 패키지의 DB 연결은 필요 없음 (나중에 수정 예정) fileType, err := option.ParseString(*n.MustGetOpt("type"), n) if err != nil { return err @@ -172,15 +170,7 @@ func ExportData(n *commando.Node) error { var data []databases.Text - conn, err := mongo.Connect(options.Client().ApplyURI(configs.Config.DatabaseURL)) - if err != nil { - ch <- err - return - } - - defer conn.Disconnect(context.TODO()) - - cur, err := conn.Database(configs.Config.DBName).Collection("text").Find(context.TODO(), bson.D{{Key: "persona", Value: "muffin"}}) + cur, err := databases.Database.Texts.Find(context.TODO(), bson.D{{Key: "persona", Value: "muffin"}}) if err != nil { ch <- err return @@ -234,15 +224,7 @@ func ExportData(n *commando.Node) error { var data []databases.Text - conn, err := mongo.Connect(options.Client().ApplyURI(configs.Config.DatabaseURL)) - if err != nil { - ch <- err - return - } - - defer conn.Disconnect(context.TODO()) - - cur, err := conn.Database(configs.Config.DBName).Collection("text").Find(context.TODO(), bson.D{ + cur, err := databases.Database.Texts.Find(context.TODO(), bson.D{ { Key: "persona", Value: bson.M{ @@ -292,15 +274,7 @@ func ExportData(n *commando.Node) error { var data []databases.Learn - conn, err := mongo.Connect(options.Client().ApplyURI(configs.Config.DatabaseURL)) - if err != nil { - ch <- err - return - } - - defer conn.Disconnect(context.TODO()) - - cur, err := conn.Database(configs.Config.DBName).Collection("learn").Find(context.TODO(), bson.D{{}}) + cur, err := databases.Database.Learns.Find(context.TODO(), bson.D{{}}) if err != nil { ch <- err return