feat: 导出密钥功能
This commit is contained in:
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"gpt-load/internal/keypool"
|
||||
"gpt-load/internal/models"
|
||||
"io"
|
||||
"regexp"
|
||||
"strings"
|
||||
|
||||
@@ -306,3 +307,28 @@ func (s *KeyService) TestMultipleKeys(group *models.Group, keysText string) ([]k
|
||||
|
||||
return allResults, nil
|
||||
}
|
||||
|
||||
// StreamKeysToWriter fetches keys from the database in batches and writes them to the provided writer.
|
||||
func (s *KeyService) StreamKeysToWriter(groupID uint, statusFilter string, writer io.Writer) error {
|
||||
query := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Select("id, key_value")
|
||||
|
||||
switch statusFilter {
|
||||
case models.KeyStatusActive, models.KeyStatusInvalid:
|
||||
query = query.Where("status = ?", statusFilter)
|
||||
case "all":
|
||||
default:
|
||||
return fmt.Errorf("invalid status filter: %s", statusFilter)
|
||||
}
|
||||
|
||||
var keys []models.APIKey
|
||||
err := query.FindInBatches(&keys, chunkSize, func(tx *gorm.DB, batch int) error {
|
||||
for _, key := range keys {
|
||||
if _, err := writer.Write([]byte(key.KeyValue + "\n")); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}).Error
|
||||
|
||||
return err
|
||||
}
|
||||
|
121
internal/services/log_service.go
Normal file
121
internal/services/log_service.go
Normal file
@@ -0,0 +1,121 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"gpt-load/internal/models"
|
||||
"io"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ExportableLogKey defines the structure for the data to be exported to CSV.
|
||||
type ExportableLogKey struct {
|
||||
KeyValue string `gorm:"column:key_value"`
|
||||
GroupName string `gorm:"column:group_name"`
|
||||
StatusCode int `gorm:"column:status_code"`
|
||||
}
|
||||
|
||||
// LogService provides services related to request logs.
|
||||
type LogService struct {
|
||||
DB *gorm.DB
|
||||
}
|
||||
|
||||
// NewLogService creates a new LogService.
|
||||
func NewLogService(db *gorm.DB) *LogService {
|
||||
return &LogService{DB: db}
|
||||
}
|
||||
|
||||
// logFiltersScope returns a GORM scope function that applies filters from the Gin context.
|
||||
func logFiltersScope(c *gin.Context) func(db *gorm.DB) *gorm.DB {
|
||||
return func(db *gorm.DB) *gorm.DB {
|
||||
if groupName := c.Query("group_name"); groupName != "" {
|
||||
db = db.Where("group_name LIKE ?", "%"+groupName+"%")
|
||||
}
|
||||
if keyValue := c.Query("key_value"); keyValue != "" {
|
||||
// 安全地处理 keyValue,避免越界错误
|
||||
var likePattern string
|
||||
if len(keyValue) > 2 {
|
||||
likePattern = "%" + keyValue[1:len(keyValue)-1] + "%"
|
||||
} else {
|
||||
likePattern = "%" + keyValue + "%"
|
||||
}
|
||||
db = db.Where("key_value LIKE ?", likePattern)
|
||||
}
|
||||
if isSuccessStr := c.Query("is_success"); isSuccessStr != "" {
|
||||
if isSuccess, err := strconv.ParseBool(isSuccessStr); err == nil {
|
||||
db = db.Where("is_success = ?", isSuccess)
|
||||
}
|
||||
}
|
||||
if statusCodeStr := c.Query("status_code"); statusCodeStr != "" {
|
||||
if statusCode, err := strconv.Atoi(statusCodeStr); err == nil {
|
||||
db = db.Where("status_code = ?", statusCode)
|
||||
}
|
||||
}
|
||||
if sourceIP := c.Query("source_ip"); sourceIP != "" {
|
||||
db = db.Where("source_ip = ?", sourceIP)
|
||||
}
|
||||
if errorContains := c.Query("error_contains"); errorContains != "" {
|
||||
db = db.Where("error_message LIKE ?", "%"+errorContains+"%")
|
||||
}
|
||||
if startTimeStr := c.Query("start_time"); startTimeStr != "" {
|
||||
if startTime, err := time.Parse(time.RFC3339, startTimeStr); err == nil {
|
||||
db = db.Where("timestamp >= ?", startTime)
|
||||
}
|
||||
}
|
||||
if endTimeStr := c.Query("end_time"); endTimeStr != "" {
|
||||
if endTime, err := time.Parse(time.RFC3339, endTimeStr); err == nil {
|
||||
db = db.Where("timestamp <= ?", endTime)
|
||||
}
|
||||
}
|
||||
return db
|
||||
}
|
||||
}
|
||||
|
||||
// GetLogsQuery returns a GORM query for fetching logs with filters.
|
||||
func (s *LogService) GetLogsQuery(c *gin.Context) *gorm.DB {
|
||||
return s.DB.Model(&models.RequestLog{}).Scopes(logFiltersScope(c))
|
||||
}
|
||||
|
||||
// StreamLogKeysToCSV fetches unique keys from logs based on filters and streams them as a CSV.
|
||||
func (s *LogService) StreamLogKeysToCSV(c *gin.Context, writer io.Writer) error {
|
||||
// Create a CSV writer
|
||||
csvWriter := csv.NewWriter(writer)
|
||||
defer csvWriter.Flush()
|
||||
|
||||
// Write CSV header
|
||||
header := []string{"key_value", "group_name", "status_code"}
|
||||
if err := csvWriter.Write(header); err != nil {
|
||||
return fmt.Errorf("failed to write CSV header: %w", err)
|
||||
}
|
||||
|
||||
// 直接查询所有结果,不使用分批,避免复杂的SQL问题
|
||||
var results []ExportableLogKey
|
||||
err := s.DB.Model(&models.RequestLog{}).
|
||||
Scopes(logFiltersScope(c)).
|
||||
Select("key_value, group_name, status_code").
|
||||
Group("key_value, group_name, status_code").
|
||||
Order("key_value").
|
||||
Find(&results).Error
|
||||
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to fetch log keys: %w", err)
|
||||
}
|
||||
|
||||
// 写入CSV数据
|
||||
for _, record := range results {
|
||||
csvRecord := []string{
|
||||
record.KeyValue,
|
||||
record.GroupName,
|
||||
strconv.Itoa(record.StatusCode),
|
||||
}
|
||||
if err := csvWriter.Write(csvRecord); err != nil {
|
||||
return fmt.Errorf("failed to write CSV record: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
Reference in New Issue
Block a user