优化

internal/data/permission.go

@@ -2,6 +2,7 @@ package data
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"kra/internal/biz"
@@ -26,9 +27,9 @@ type menuButtonPO struct {
 func (menuButtonPO) TableName() string { return "sys_base_menu_btns" }
 
 type authorityButtonPO struct {
-	AuthorityID uint `gorm:"primaryKey;column:authority_id"`
-	MenuID      uint `gorm:"primaryKey;column:sys_menu_id"`
-	ButtonID    uint `gorm:"primaryKey;column:sys_base_menu_btn_id"`
+	AuthorityID uint `gorm:"column:authority_id"`
+	MenuID      uint `gorm:"column:sys_menu_id"`
+	ButtonID    uint `gorm:"column:sys_base_menu_btn_id"`
 }
 
 func (authorityButtonPO) TableName() string { return "sys_authority_btns" }
@@ -62,9 +63,15 @@ func (r *permissionRepo) SetAuthorityButtons(ctx context.Context, aid uint, butt
 	})
 }
 func (r *permissionRepo) SelectedButtons(ctx context.Context, aid, menuID uint) ([]uint, error) {
+	var rows []authorityButtonPO
+	if err := r.data.gormDB.WithContext(ctx).Where("authority_id = ? AND sys_menu_id = ?", aid, menuID).Find(&rows).Error; err != nil {
+		return nil, err
+	}
 	var ids []uint
-	err := r.data.gormDB.WithContext(ctx).Model(&authorityButtonPO{}).Where("authority_id = ? AND sys_menu_id = ?", aid, menuID).Pluck("sys_base_menu_btn_id", &ids).Error
-	return ids, err
+	for _, row := range rows {
+		ids = append(ids, row.ButtonID)
+	}
+	return ids, nil
 }
 func (r *permissionRepo) SetSelectedButtons(ctx context.Context, aid, menuID uint, ids []uint) error {
 	return r.data.gormDB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
@@ -81,10 +88,16 @@ func (r *permissionRepo) SetSelectedButtons(ctx context.Context, aid, menuID uin
 		return nil
 	})
 }
-func (r *permissionRepo) CanRemoveButton(ctx context.Context, id uint) (bool, error) {
-	var count int64
-	err := r.data.gormDB.WithContext(ctx).Model(&authorityButtonPO{}).Where("sys_base_menu_btn_id = ?", id).Count(&count).Error
-	return count == 0, err
+func (r *permissionRepo) CanRemoveButton(ctx context.Context, id string) (bool, error) {
+	var row authorityButtonPO
+	err := r.data.gormDB.WithContext(ctx).Where("sys_base_menu_btn_id = ?", id).First(&row).Error
+	if errors.Is(err, gorm.ErrRecordNotFound) {
+		return true, nil
+	}
+	if err != nil {
+		return false, errors.New("此按钮正在被使用无法删除")
+	}
+	return false, nil
 }
 func (r *permissionRepo) AuthorityButtonIDs(ctx context.Context, aid uint) ([]uint, error) {
 	var ids []uint