Enterprise-Knowledge-Base/internal/application/repository/task_queue.go

package repository

import (
	"context"
	"errors"
	"fmt"
	"strconv"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"gorm.io/gorm"
)

// taskPendingOpsRepository implements interfaces.TaskPendingOpsRepository.
type taskPendingOpsRepository struct {
	db *gorm.DB
}

// NewTaskPendingOpsRepository constructs a GORM-backed implementation.
func NewTaskPendingOpsRepository(db *gorm.DB) interfaces.TaskPendingOpsRepository {
	return &taskPendingOpsRepository{db: db}
}

// Enqueue inserts a single op. Callers must populate TenantID/TaskType/
// Scope/ScopeID/Op (Payload optional). ID, FailCount default to zero;
// EnqueuedAt is filled with the DB-side default if left zero.
func (r *taskPendingOpsRepository) Enqueue(ctx context.Context, op *types.TaskPendingOp) error {
	if op == nil {
		return errors.New("task pending ops: nil op")
	}
	if op.TaskType == "" || op.Scope == "" || op.ScopeID == "" {
		return errors.New("task pending ops: task_type, scope, scope_id are required")
	}
	if op.Op == "" {
		return errors.New("task pending ops: op is required")
	}
	if len(op.Payload) == 0 {
		// Make sure the JSONB column never sees NULL — the schema sets a
		// default but explicit "{}" keeps the row uniform regardless of
		// driver-level default handling.
		op.Payload = []byte("{}")
	}
	return r.db.WithContext(ctx).Create(op).Error
}

// PeekBatch returns up to `limit` rows for the (task_type, scope, scope_id)
// tuple ordered by id ASC. Rows are not removed; callers must
// DeleteByIDs once they have been consumed (or IncrFailCount and leave
// them for the next pass). `limit` <= 0 falls back to 1; we clamp the
// upper bound generously so callers can pull large windows when they
// know the consumer can handle them.
func (r *taskPendingOpsRepository) PeekBatch(
	ctx context.Context,
	taskType, scope, scopeID string,
	limit int,
) ([]*types.TaskPendingOp, error) {
	if limit <= 0 {
		limit = 1
	}
	if limit > 1000 {
		limit = 1000
	}
	var ops []*types.TaskPendingOp
	if err := r.db.WithContext(ctx).
		Where("task_type = ? AND scope = ? AND scope_id = ?", taskType, scope, scopeID).
		Order("id ASC").
		Limit(limit).
		Find(&ops).Error; err != nil {
		return nil, err
	}
	return ops, nil
}

// DeleteByIDs removes the given rows in one statement. Empty input is a
// no-op so the caller can invoke unconditionally at the end of a batch.
func (r *taskPendingOpsRepository) DeleteByIDs(ctx context.Context, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}
	return r.db.WithContext(ctx).
		Where("id IN ?", ids).
		Delete(&types.TaskPendingOp{}).Error
}

// IncrFailCount atomically bumps fail_count for one row and returns the
// new value. We use UPDATE ... RETURNING so the read+write happens in
// one round trip and races between concurrent IncrFailCount callers
// resolve to monotonic counts.
//
// A missing row returns (0, nil): the caller's ID may have been removed
// by a concurrent DeleteByIDs (e.g. dead-letter path), which is benign.
func (r *taskPendingOpsRepository) IncrFailCount(ctx context.Context, id int64) (int, error) {
	var newCount int
	err := r.db.WithContext(ctx).Raw(
		`UPDATE task_pending_ops SET fail_count = fail_count + 1 WHERE id = ? RETURNING fail_count`,
		id,
	).Scan(&newCount).Error
	if err != nil {
		return 0, err
	}
	return newCount, nil
}

// PendingCount returns how many rows are currently queued for the
// tuple. Covered by idx_task_pending_ops_scope.
func (r *taskPendingOpsRepository) PendingCount(
	ctx context.Context,
	taskType, scope, scopeID string,
) (int64, error) {
	var n int64
	if err := r.db.WithContext(ctx).
		Model(&types.TaskPendingOp{}).
		Where("task_type = ? AND scope = ? AND scope_id = ?", taskType, scope, scopeID).
		Count(&n).Error; err != nil {
		return 0, err
	}
	return n, nil
}

// DeleteByDedupKey drops rows in the tuple whose dedup_key matches.
// If `op` is non-empty, only rows with the matching op are dropped;
// otherwise every matching row is removed. Empty dedup_key is rejected
// to prevent accidentally wiping the entire queue for a KB.
//
// Used by:
//   - Wiki delete path: scrub queued WikiOpIngest entries for a
//     knowledge that is being deleted, while preserving WikiOpRetract
//     so the cleanup can still unlink pages.
//   - Wiki reparse path: same scrub of pending ingests so the new
//     parse can repopulate cleanly.
func (r *taskPendingOpsRepository) DeleteByDedupKey(
	ctx context.Context,
	taskType, scope, scopeID, dedupKey, op string,
) error {
	if dedupKey == "" {
		return fmt.Errorf("task pending ops: empty dedup_key in DeleteByDedupKey")
	}
	q := r.db.WithContext(ctx).
		Where("task_type = ? AND scope = ? AND scope_id = ? AND dedup_key = ?",
			taskType, scope, scopeID, dedupKey)
	if op != "" {
		q = q.Where("op = ?", op)
	}
	return q.Delete(&types.TaskPendingOp{}).Error
}

// taskDeadLetterRepository implements interfaces.TaskDeadLetterRepository.
type taskDeadLetterRepository struct {
	db *gorm.DB
}

// NewTaskDeadLetterRepository constructs a GORM-backed implementation.
func NewTaskDeadLetterRepository(db *gorm.DB) interfaces.TaskDeadLetterRepository {
	return &taskDeadLetterRepository{db: db}
}

// Insert records one dead letter. Best-effort caller: the asynq
// middleware swallows the error so a failed insert never masks the
// underlying task error.
func (r *taskDeadLetterRepository) Insert(ctx context.Context, dl *types.TaskDeadLetter) error {
	if dl == nil {
		return errors.New("task dead letters: nil entry")
	}
	if dl.TaskType == "" {
		return errors.New("task dead letters: task_type is required")
	}
	if dl.Scope == "" {
		dl.Scope = types.TaskScopeUnknown
	}
	if len(dl.Payload) == 0 {
		dl.Payload = []byte("{}")
	}
	return r.db.WithContext(ctx).Create(dl).Error
}

// ListByScope returns dead letters for (scope, scope_id) newest-first
// with a stringified id cursor. `limit` is clamped to [1, 200]. Empty
// nextCursor signals the tail.
func (r *taskDeadLetterRepository) ListByScope(
	ctx context.Context,
	scope, scopeID, cursor string,
	limit int,
) ([]*types.TaskDeadLetter, string, error) {
	if scope == "" || scopeID == "" {
		return nil, "", errors.New("task dead letters: scope and scope_id are required")
	}
	return r.list(ctx, cursor, limit, func(q *gorm.DB) *gorm.DB {
		return q.Where("scope = ? AND scope_id = ?", scope, scopeID)
	})
}

// ListByTaskType returns dead letters for the given task_type
// newest-first with a stringified id cursor. Same clamping rules.
func (r *taskDeadLetterRepository) ListByTaskType(
	ctx context.Context,
	taskType, cursor string,
	limit int,
) ([]*types.TaskDeadLetter, string, error) {
	if taskType == "" {
		return nil, "", errors.New("task dead letters: task_type is required")
	}
	return r.list(ctx, cursor, limit, func(q *gorm.DB) *gorm.DB {
		return q.Where("task_type = ?", taskType)
	})
}

// list is the shared cursor pagination implementation, parametrized by
// the caller-supplied filter. Mirrors wikiLogEntryRepository.List.
func (r *taskDeadLetterRepository) list(
	ctx context.Context,
	cursor string,
	limit int,
	filter func(*gorm.DB) *gorm.DB,
) ([]*types.TaskDeadLetter, string, error) {
	if limit <= 0 {
		limit = 50
	}
	if limit > 200 {
		limit = 200
	}

	q := r.db.WithContext(ctx).Order("id DESC").Limit(limit)
	q = filter(q)

	if cursor != "" {
		cursorID, err := strconv.ParseInt(cursor, 10, 64)
		if err != nil {
			return nil, "", fmt.Errorf("invalid cursor %q: %w", cursor, err)
		}
		q = q.Where("id < ?", cursorID)
	}

	var rows []*types.TaskDeadLetter
	if err := q.Find(&rows).Error; err != nil {
		return nil, "", err
	}

	nextCursor := ""
	if len(rows) == limit {
		nextCursor = strconv.FormatInt(rows[len(rows)-1].ID, 10)
	}
	return rows, nextCursor, nil
}

// DeleteByID drops a single dead letter row. Returns nil even if the
// row is already gone — operators issuing concurrent deletes shouldn't
// see spurious errors.
func (r *taskDeadLetterRepository) DeleteByID(ctx context.Context, id int64) error {
	return r.db.WithContext(ctx).
		Where("id = ?", id).
		Delete(&types.TaskDeadLetter{}).Error
}