39 lines
926 B
Go
39 lines
926 B
Go
|
package database
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
|
||
|
|
"gorm.io/gorm"
|
||
|
|
)
|
||
|
|
|
||
|
|
// GetTableName 获取带前缀的表名
|
||
|
|
// 格式：{botID}_{tableName}
|
||
|
|
func GetTableName(botID, tableName string) string {
|
||
|
|
if botID == "" {
|
||
|
|
return tableName
|
||
|
|
}
|
||
|
|
// 清理 botID 中的特殊字符，确保表名合法
|
||
|
|
cleanBotID := sanitizeTableName(botID)
|
||
|
|
return fmt.Sprintf("%s_%s", cleanBotID, tableName)
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithTablePrefix 为查询添加表前缀
|
||
|
|
func WithTablePrefix(db *gorm.DB, botID, tableName string) *gorm.DB {
|
||
|
|
prefixedTableName := GetTableName(botID, tableName)
|
||
|
|
return db.Table(prefixedTableName)
|
||
|
|
}
|
||
|
|
|
||
|
|
// sanitizeTableName 清理表名中的特殊字符
|
||
|
|
func sanitizeTableName(name string) string {
|
||
|
|
// 移除或替换 SQL 不安全的字符
|
||
|
|
result := ""
|
||
|
|
for _, r := range name {
|
||
|
|
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '_' {
|
||
|
|
result += string(r)
|
||
|
|
} else {
|
||
|
|
result += "_"
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return result
|
||
|
|
}
|