feat: set dir for storages

dao/callback_data.go

@@ -0,0 +1,19 @@
+package dao
+
+func CreateCallbackData(data string) (uint, error) {
+	callbackData := CallbackData{
+		Data: data,
+	}
+	err := db.Create(&callbackData).Error
+	return callbackData.ID, err
+}
+
+func GetCallbackData(id uint) (string, error) {
+	var callbackData CallbackData
+	err := db.First(&callbackData, id).Error
+	return callbackData.Data, err
+}
+
+func DeleteCallbackData(id uint) error {
+	return db.Unscoped().Where("id = ?", id).Delete(&CallbackData{}).Error
+}

dao/db.go

@@ -9,7 +9,6 @@ import (
 	"github.com/glebarez/sqlite"
 	"github.com/krau/SaveAny-Bot/config"
 	"github.com/krau/SaveAny-Bot/logger"
-	"github.com/krau/SaveAny-Bot/types"
 	"gorm.io/gorm"
 	glogger "gorm.io/gorm/logger"
 )
@@ -37,7 +36,7 @@ func Init() {
 		os.Exit(1)
 	}
 	logger.L.Debug("Database connected")
-	if err := db.AutoMigrate(&types.ReceivedFile{}, &types.User{}); err != nil {
+	if err := db.AutoMigrate(&ReceivedFile{}, &User{}, &Dir{}, &CallbackData{}); err != nil {
 		logger.L.Fatal("迁移数据库失败, 如果您从旧版本升级, 建议手动删除数据库文件后重试: ", err)
 	}
 
@@ -52,7 +51,7 @@ func syncUsers() error {
 		return fmt.Errorf("failed to get users: %w", err)
 	}
 
-	dbUserMap := make(map[int64]types.User)
+	dbUserMap := make(map[int64]User)
 	for _, u := range dbUsers {
 		dbUserMap[u.ChatID] = u
 	}

dao/dir.go

@@ -0,0 +1,43 @@
+package dao
+
+func CreateDirForUser(userID uint, storageName, path string) error {
+	dir := Dir{
+		UserID:      userID,
+		StorageName: storageName,
+		Path:        path,
+	}
+	return db.Create(&dir).Error
+}
+
+func GetDirByID(id uint) (*Dir, error) {
+	dir := &Dir{}
+	err := db.First(dir, id).Error
+	if err != nil {
+		return nil, err
+	}
+	return dir, err
+}
+
+func GetUserDirs(userID uint) ([]Dir, error) {
+	var dirs []Dir
+	err := db.Where("user_id = ?", userID).Find(&dirs).Error
+	return dirs, err
+}
+
+func GetUserDirsByChatID(chatID int64) ([]Dir, error) {
+	user, err := GetUserByChatID(chatID)
+	if err != nil {
+		return nil, err
+	}
+	return GetUserDirs(user.ID)
+}
+
+func GetDirsByUserIDAndStorageName(userID uint, storageName string) ([]Dir, error) {
+	var dirs []Dir
+	err := db.Where("user_id = ? AND storage_name = ?", userID, storageName).Find(&dirs).Error
+	return dirs, err
+}
+
+func DeleteDirForUser(userID uint, storageName, path string) error {
+	return db.Unscoped().Where("user_id = ? AND storage_name = ? AND path = ?", userID, storageName, path).Delete(&Dir{}).Error
+}

dao/file.go

@@ -1,8 +1,6 @@
 package dao
 
-import "github.com/krau/SaveAny-Bot/types"
-
-func SaveReceivedFile(receivedFile *types.ReceivedFile) error {
+func SaveReceivedFile(receivedFile *ReceivedFile) error {
 	record, err := GetReceivedFileByChatAndMessageID(receivedFile.ChatID, receivedFile.MessageID)
 	if err == nil {
 		receivedFile.ID = record.ID
@@ -10,8 +8,8 @@ func SaveReceivedFile(receivedFile *types.ReceivedFile) error {
 	return db.Save(receivedFile).Error
 }
 
-func GetReceivedFileByChatAndMessageID(chatID int64, messageID int) (*types.ReceivedFile, error) {
-	var receivedFile types.ReceivedFile
+func GetReceivedFileByChatAndMessageID(chatID int64, messageID int) (*ReceivedFile, error) {
+	var receivedFile ReceivedFile
 	err := db.Where("chat_id = ? AND message_id = ?", chatID, messageID).First(&receivedFile).Error
 	if err != nil {
 		return nil, err
@@ -19,6 +17,6 @@ func GetReceivedFileByChatAndMessageID(chatID int64, messageID int) (*types.Rece
 	return &receivedFile, nil
 }
 
-func DeleteReceivedFile(receivedFile *types.ReceivedFile) error {
+func DeleteReceivedFile(receivedFile *ReceivedFile) error {
 	return db.Unscoped().Delete(receivedFile).Error
 }

dao/model.go

@@ -0,0 +1,37 @@
+package dao
+
+import (
+	"gorm.io/gorm"
+)
+
+type ReceivedFile struct {
+	gorm.Model
+	Processing bool
+	// Which chat the file is from
+	ChatID int64 `gorm:"uniqueIndex:idx_chat_id_message_id;not null"`
+	// Which message the file is from
+	MessageID      int `gorm:"uniqueIndex:idx_chat_id_message_id;not null"`
+	ReplyMessageID int
+	ReplyChatID    int64
+	FileName       string
+}
+
+type User struct {
+	gorm.Model
+	ChatID         int64 `gorm:"uniqueIndex;not null"`
+	Silent         bool
+	DefaultStorage string // Default storage name
+	Dirs           []Dir
+}
+
+type Dir struct {
+	gorm.Model
+	UserID      uint
+	StorageName string
+	Path        string
+}
+
+type CallbackData struct {
+	gorm.Model
+	Data string
+}

dao/user.go

@@ -1,32 +1,30 @@
 package dao
 
-import (
-	"github.com/krau/SaveAny-Bot/types"
-)
-
 func CreateUser(chatID int64) error {
 	if _, err := GetUserByChatID(chatID); err == nil {
 		return nil
 	}
-	return db.Create(&types.User{ChatID: chatID}).Error
+	return db.Create(&User{ChatID: chatID}).Error
 }
 
-func GetAllUsers() ([]types.User, error) {
-	var users []types.User
-	err := db.Find(&users).Error
+func GetAllUsers() ([]User, error) {
+	var users []User
+	err := db.Preload("Dirs").Find(&users).Error
 	return users, err
 }
 
-func GetUserByChatID(chatID int64) (*types.User, error) {
-	var user types.User
-	err := db.Where("chat_id = ?", chatID).First(&user).Error
+func GetUserByChatID(chatID int64) (*User, error) {
+	var user User
+	err := db.
+		Preload("Dirs").
+		Where("chat_id = ?", chatID).First(&user).Error
 	return &user, err
 }
 
-func UpdateUser(user *types.User) error {
+func UpdateUser(user *User) error {
 	return db.Save(user).Error
 }
 
-func DeleteUser(user *types.User) error {
-	return db.Unscoped().Delete(user).Error
+func DeleteUser(user *User) error {
+	return db.Unscoped().Select("Dirs").Delete(user).Error
 }