1 parent
38c4f63e2e
commit
4721c1e97f
11 files changed
+389
-3
No files matched your search
@@ -58,6 +58,9 @@ func (r *Repository) BulkDisableUsers(actorID string, userIDs []string, events [
|
||||
if err := tx.Delete(&model.AuthSession{}, "user_id IN ?", userIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Delete(&model.TaskTextDelta{}, "user_id IN ?", userIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Model(&model.User{}).Where("id IN ?", userIDs).Updates(map[string]any{"status": model.UserStatusDisabled, "updated_at": now}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -222,7 +222,7 @@ func (r *Repository) RetryTaskWithBilling(userID string, taskID string, order *m
|
||||
}
|
||||
updates := map[string]any{
|
||||
"status": model.TaskStatusQueued, "stage": "等待队列调度", "progress": 5, "error": "", "result_json": "",
|
||||
"started_at": nil, "completed_at": nil, "updated_at": time.Now(),
|
||||
"text_draft": "", "started_at": nil, "completed_at": nil, "updated_at": time.Now(),
|
||||
}
|
||||
if order != nil {
|
||||
updates["billing_order_id"] = order.ID
|
||||
@@ -239,6 +239,9 @@ func (r *Repository) RetryTaskWithBilling(userID string, taskID string, order *m
|
||||
if updated.RowsAffected != 1 {
|
||||
return ErrTaskNotRetryable
|
||||
}
|
||||
if err := tx.Delete(&model.TaskTextDelta{}, "user_id = ? AND task_id = ?", userID, taskID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.First(&task, "id = ? AND user_id = ?", taskID, userID).Error
|
||||
})
|
||||
return &task, err
|
||||
|
||||
@@ -16,6 +16,10 @@ var ErrDailyUploadLimitExceeded = errors.New("daily upload limit exceeded")
|
||||
|
||||
var ErrTaskProviderRecoveryConflict = errors.New("task provider recovery is already running")
|
||||
|
||||
var ErrTextReplayQuotaExceeded = errors.New("text replay quota exceeded")
|
||||
|
||||
var ErrTextReplayClosed = errors.New("text replay task is closed")
|
||||
|
||||
type Repository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
@@ -52,9 +56,10 @@ func (r *Repository) UserStorageUsage(userID string) (UserStorageUsage, error) {
|
||||
(SELECT COALESCE(SUM(length(CAST(COALESCE(prompt, '') AS BLOB)) + length(CAST(COALESCE(canvas_snapshot_json, '') AS BLOB)) + length(CAST(COALESCE(canvas_ops_json, '') AS BLOB))), 0) FROM sessions WHERE user_id = ?)
|
||||
+ (SELECT COALESCE(SUM(length(CAST(COALESCE(content, '') AS BLOB)) + length(CAST(COALESCE(payload, '') AS BLOB))), 0) FROM messages WHERE user_id = ?) AS session_bytes,
|
||||
(SELECT COUNT(*) FROM tasks WHERE user_id = ?) AS task_count,
|
||||
(SELECT COALESCE(SUM(length(CAST(COALESCE(prompt, '') AS BLOB)) + length(CAST(COALESCE(input_json, '') AS BLOB)) + length(CAST(COALESCE(result_json, '') AS BLOB)) + length(CAST(COALESCE(error, '') AS BLOB))), 0) FROM tasks WHERE user_id = ?)
|
||||
(SELECT COALESCE(SUM(length(CAST(COALESCE(prompt, '') AS BLOB)) + length(CAST(COALESCE(input_json, '') AS BLOB)) + length(CAST(COALESCE(result_json, '') AS BLOB)) + length(CAST(COALESCE(text_draft, '') AS BLOB)) + length(CAST(COALESCE(error, '') AS BLOB))), 0) FROM tasks WHERE user_id = ?)
|
||||
+ (SELECT COALESCE(SUM(length(CAST(COALESCE(message, '') AS BLOB)) + length(CAST(COALESCE(payload, '') AS BLOB))), 0) FROM task_logs WHERE user_id = ?)
|
||||
+ (SELECT COALESCE(SUM(length(CAST(COALESCE(url, '') AS BLOB)) + length(CAST(COALESCE(payload, '') AS BLOB))), 0) FROM results WHERE user_id = ?)
|
||||
+ (SELECT COALESCE(SUM(byte_count), 0) FROM task_text_deltas WHERE user_id = ?)
|
||||
+ (SELECT COALESCE(SUM(length(CAST(COALESCE(path, '') AS BLOB)) + length(CAST(COALESCE(model, '') AS BLOB)) + length(CAST(COALESCE(provider_request_id, '') AS BLOB)) + length(CAST(COALESCE(error_code, '') AS BLOB)) + length(CAST(COALESCE(error, '') AS BLOB)) + length(CAST(COALESCE(upstream_url, '') AS BLOB)) + length(CAST(COALESCE(request_body, '') AS BLOB)) + length(CAST(COALESCE(response_body, '') AS BLOB))), 0) FROM api_call_logs WHERE user_id = ?) AS task_bytes,
|
||||
(SELECT COUNT(*) FROM api_call_logs WHERE user_id = ?) AS api_call_count
|
||||
`
|
||||
@@ -62,7 +67,7 @@ func (r *Repository) UserStorageUsage(userID string) (UserStorageUsage, error) {
|
||||
query = strings.ReplaceAll(query, "length(CAST(COALESCE(", "octet_length(COALESCE(")
|
||||
query = strings.ReplaceAll(query, ", '') AS BLOB))", ", ''))")
|
||||
}
|
||||
err := r.db.Raw(query, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID).Scan(&usage).Error
|
||||
err := r.db.Raw(query, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID, userID).Scan(&usage).Error
|
||||
return usage, err
|
||||
}
|
||||
|
||||
@@ -450,6 +455,15 @@ func (r *Repository) SessionForUser(userID string, id string) (*model.Session, e
|
||||
|
||||
func (r *Repository) DeleteSessionDraft(userID string, id string) error {
|
||||
return r.db.Transaction(func(tx *gorm.DB) error {
|
||||
var taskIDs []string
|
||||
if err := tx.Model(&model.Task{}).Where("user_id = ? AND session_id = ?", userID, id).Pluck("id", &taskIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(taskIDs) > 0 {
|
||||
if err := tx.Delete(&model.TaskTextDelta{}, "user_id = ? AND task_id IN ?", userID, taskIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := tx.Delete(&model.Message{}, "user_id = ? AND session_id = ?", userID, id).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"infinite-canvas/backend/internal/model"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type TextReplayLimits struct {
|
||||
MaxTaskBytes int64
|
||||
MaxUserBytes int64
|
||||
MaxTaskEvents int64
|
||||
}
|
||||
|
||||
type TextReplayStats struct {
|
||||
EventCount int64 `json:"eventCount"`
|
||||
TaskCount int64 `json:"taskCount"`
|
||||
ByteCount int64 `json:"byteCount"`
|
||||
OldestAt *time.Time `json:"oldestAt,omitempty"`
|
||||
}
|
||||
|
||||
func (r *Repository) AppendTaskTextDelta(userID string, taskID string, content string, expiresAt time.Time, limits TextReplayLimits) (*model.TaskTextDelta, error) {
|
||||
var created model.TaskTextDelta
|
||||
err := r.db.Transaction(func(tx *gorm.DB) error {
|
||||
var task model.Task
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id", "user_id", "status").First(&task, "id = ? AND user_id = ?", taskID, userID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.Status == model.TaskStatusSucceeded || task.Status == model.TaskStatusFailed || task.Status == model.TaskStatusCancelled {
|
||||
return ErrTextReplayClosed
|
||||
}
|
||||
var taskUsage struct {
|
||||
Count int64
|
||||
Bytes int64
|
||||
Max int64
|
||||
}
|
||||
if err := tx.Model(&model.TaskTextDelta{}).Select("count(*) AS count, COALESCE(sum(byte_count), 0) AS bytes, COALESCE(max(sequence), 0) AS max").Where("task_id = ?", taskID).Scan(&taskUsage).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var userBytes int64
|
||||
if err := tx.Model(&model.TaskTextDelta{}).Select("COALESCE(sum(byte_count), 0)").Where("user_id = ?", userID).Scan(&userBytes).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
incoming := int64(len([]byte(content)))
|
||||
if taskUsage.Count+1 > limits.MaxTaskEvents || taskUsage.Bytes+incoming > limits.MaxTaskBytes || userBytes+incoming > limits.MaxUserBytes {
|
||||
return ErrTextReplayQuotaExceeded
|
||||
}
|
||||
created = model.TaskTextDelta{ID: newRepositoryID(), UserID: userID, TaskID: taskID, Sequence: taskUsage.Max + 1, Content: content, ByteCount: incoming, CreatedAt: time.Now(), ExpiresAt: expiresAt}
|
||||
return tx.Create(&created).Error
|
||||
})
|
||||
return &created, err
|
||||
}
|
||||
|
||||
func (r *Repository) TaskTextDeltas(userID string, taskID string, after int64, limit int) ([]model.TaskTextDelta, error) {
|
||||
var items []model.TaskTextDelta
|
||||
if limit <= 0 || limit > 1000 {
|
||||
limit = 1000
|
||||
}
|
||||
err := r.db.Where("user_id = ? AND task_id = ? AND sequence > ?", userID, taskID, after).Order("sequence asc").Limit(limit).Find(&items).Error
|
||||
return items, err
|
||||
}
|
||||
|
||||
func (r *Repository) CompactTaskTextDeltas(taskID string, expiresAt time.Time, keepDraft bool) error {
|
||||
return r.db.Transaction(func(tx *gorm.DB) error {
|
||||
var task model.Task
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&task, "id = ?", taskID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if keepDraft {
|
||||
var items []model.TaskTextDelta
|
||||
if err := tx.Select("content").Where("task_id = ?", taskID).Order("sequence asc").Find(&items).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if task.TextDraft == "" {
|
||||
for _, item := range items {
|
||||
task.TextDraft += item.Content
|
||||
}
|
||||
if err := tx.Model(&task).Update("text_draft", task.TextDraft).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return tx.Model(&model.TaskTextDelta{}).Where("task_id = ?", taskID).Update("expires_at", expiresAt).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *Repository) CleanupTaskTextDeltas(now time.Time) (int64, error) {
|
||||
var deleted int64
|
||||
err := r.db.Transaction(func(tx *gorm.DB) error {
|
||||
var taskIDs []string
|
||||
if err := tx.Model(&model.TaskTextDelta{}).Where("expires_at <= ?", now).Distinct().Pluck("task_id", &taskIDs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, taskID := range taskIDs {
|
||||
var task model.Task
|
||||
if err := tx.Select("id", "status", "text_draft").First(&task, "id = ?", taskID).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
if task.TextDraft != "" || (task.Status != model.TaskStatusFailed && task.Status != model.TaskStatusCancelled) {
|
||||
continue
|
||||
}
|
||||
var items []model.TaskTextDelta
|
||||
if err := tx.Select("content").Where("task_id = ?", taskID).Order("sequence asc").Find(&items).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, item := range items {
|
||||
task.TextDraft += item.Content
|
||||
}
|
||||
if err := tx.Model(&task).Update("text_draft", task.TextDraft).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
result := tx.Where("expires_at <= ? OR NOT EXISTS (SELECT 1 FROM tasks WHERE tasks.id = task_text_deltas.task_id)", now).Delete(&model.TaskTextDelta{})
|
||||
deleted = result.RowsAffected
|
||||
return result.Error
|
||||
})
|
||||
return deleted, err
|
||||
}
|
||||
|
||||
func (r *Repository) DeleteUserTaskTextDeltas(userID string) error {
|
||||
return r.db.Delete(&model.TaskTextDelta{}, "user_id = ?", userID).Error
|
||||
}
|
||||
|
||||
func (r *Repository) TextReplayStats() (TextReplayStats, error) {
|
||||
var stats TextReplayStats
|
||||
err := r.db.Model(&model.TaskTextDelta{}).Select("count(*) AS event_count, count(DISTINCT task_id) AS task_count, COALESCE(sum(byte_count), 0) AS byte_count, min(created_at) AS oldest_at").Scan(&stats).Error
|
||||
return stats, err
|
||||
}
|
||||
Reference in new issue
Block a user