Files
geo/server/internal/tenant/app/schedule_task_service.go
T
root 6869dc2fa6 feat(schedule): enhance schedule tasks with web search and batch generation support
Add enable_web_search and generate_count columns to schedule_tasks table.
Refactor schedule task service to publish generation jobs to RabbitMQ
instead of inline execution, supporting batch dispatch and backpressure
awareness via scheduler process.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-15 14:20:54 +08:00

465 lines
16 KiB
Go

package app
import (
"context"
"encoding/json"
"fmt"
"time"
"github.com/jackc/pgx/v5/pgtype"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/geo-platform/tenant-api/internal/shared/auditlog"
"github.com/geo-platform/tenant-api/internal/shared/auth"
"github.com/geo-platform/tenant-api/internal/shared/middleware"
"github.com/geo-platform/tenant-api/internal/shared/response"
sharedschedule "github.com/geo-platform/tenant-api/internal/shared/schedule"
)
type ScheduleTaskService struct {
pool *pgxpool.Pool
auditLogs *auditlog.AsyncWriter
}
func NewScheduleTaskService(pool *pgxpool.Pool, auditLogs *auditlog.AsyncWriter) *ScheduleTaskService {
return &ScheduleTaskService{pool: pool, auditLogs: auditLogs}
}
type ScheduleTaskRequest struct {
PromptRuleID int64 `json:"prompt_rule_id" binding:"required"`
BrandID *int64 `json:"brand_id"`
Name string `json:"name" binding:"required"`
CronExpr string `json:"cron_expr" binding:"required"`
TargetPlatform *string `json:"target_platform"`
EnableWebSearch *bool `json:"enable_web_search"`
GenerateCount *int `json:"generate_count"`
StartAt *string `json:"start_at"`
EndAt *string `json:"end_at"`
}
type ScheduleTaskResponse struct {
ID int64 `json:"id"`
PromptRuleID int64 `json:"prompt_rule_id"`
PromptRuleName *string `json:"prompt_rule_name"`
BrandID *int64 `json:"brand_id"`
BrandName *string `json:"brand_name"`
Name string `json:"name"`
CronExpr string `json:"cron_expr"`
TargetPlatform *string `json:"target_platform"`
EnableWebSearch bool `json:"enable_web_search"`
GenerateCount int `json:"generate_count"`
StartAt *string `json:"start_at"`
EndAt *string `json:"end_at"`
NextRunAt *string `json:"next_run_at"`
Status string `json:"status"`
CreatedAt string `json:"created_at"`
UpdatedAt string `json:"updated_at"`
}
type ScheduleTaskListParams struct {
PromptRuleID *int64 `json:"prompt_rule_id"`
Status *string `json:"status"`
Keyword *string `json:"keyword"`
CreatedFrom *time.Time
CreatedTo *time.Time
Page int `json:"page"`
PageSize int `json:"page_size"`
}
type ScheduleTaskListResponse struct {
Items []ScheduleTaskResponse `json:"items"`
Total int64 `json:"total"`
}
type ScheduleTaskStatusRequest struct {
Status string `json:"status" binding:"required,oneof=enabled disabled"`
}
func (s *ScheduleTaskService) List(ctx context.Context, params ScheduleTaskListParams) (*ScheduleTaskListResponse, error) {
actor := auth.MustActor(ctx)
if params.Page < 1 {
params.Page = 1
}
if params.PageSize < 1 || params.PageSize > 100 {
params.PageSize = 20
}
offset := (params.Page - 1) * params.PageSize
var total int64
err := s.pool.QueryRow(ctx, `
SELECT COUNT(*) FROM schedule_tasks
WHERE tenant_id = $1 AND deleted_at IS NULL
AND ($2::bigint IS NULL OR prompt_rule_id = $2)
AND ($3::text IS NULL OR status = $3)
AND ($4::text IS NULL OR name ILIKE '%' || $4 || '%')
AND ($5::timestamptz IS NULL OR created_at >= $5)
AND ($6::timestamptz IS NULL OR created_at <= $6)
`, actor.TenantID, params.PromptRuleID, params.Status, params.Keyword, params.CreatedFrom, params.CreatedTo).Scan(&total)
if err != nil {
return nil, response.ErrInternal(50010, "count_failed", "failed to count schedule tasks")
}
rows, err := s.pool.Query(ctx, `
SELECT st.id, st.prompt_rule_id, st.brand_id, st.name,
st.cron_expr, st.target_platform, st.enable_web_search, st.generate_count, st.start_at, st.end_at,
st.next_run_at, st.status, st.created_at, st.updated_at,
pr.name AS prompt_rule_name,
b.name AS brand_name
FROM schedule_tasks st
LEFT JOIN prompt_rules pr ON pr.id = st.prompt_rule_id
LEFT JOIN brands b ON b.id = st.brand_id
WHERE st.tenant_id = $1 AND st.deleted_at IS NULL
AND ($2::bigint IS NULL OR st.prompt_rule_id = $2)
AND ($3::text IS NULL OR st.status = $3)
AND ($4::text IS NULL OR st.name ILIKE '%' || $4 || '%')
AND ($5::timestamptz IS NULL OR st.created_at >= $5)
AND ($6::timestamptz IS NULL OR st.created_at <= $6)
ORDER BY st.created_at DESC
LIMIT $7 OFFSET $8
`, actor.TenantID, params.PromptRuleID, params.Status, params.Keyword, params.CreatedFrom, params.CreatedTo, params.PageSize, offset)
if err != nil {
return nil, response.ErrInternal(50010, "query_failed", "failed to list schedule tasks")
}
defer rows.Close()
var items []ScheduleTaskResponse
for rows.Next() {
var r ScheduleTaskResponse
var ca, ua interface{}
var startAt, endAt, nextRunAt interface{}
if err := rows.Scan(&r.ID, &r.PromptRuleID, &r.BrandID, &r.Name,
&r.CronExpr, &r.TargetPlatform, &r.EnableWebSearch, &r.GenerateCount, &startAt, &endAt,
&nextRunAt, &r.Status, &ca, &ua,
&r.PromptRuleName, &r.BrandName); err != nil {
return nil, response.ErrInternal(50010, "scan_failed", err.Error())
}
r.CreatedAt = fmt.Sprintf("%v", ca)
r.UpdatedAt = fmt.Sprintf("%v", ua)
if startAt != nil {
s := fmt.Sprintf("%v", startAt)
r.StartAt = &s
}
if endAt != nil {
s := fmt.Sprintf("%v", endAt)
r.EndAt = &s
}
if nextRunAt != nil {
s := fmt.Sprintf("%v", nextRunAt)
r.NextRunAt = &s
}
items = append(items, r)
}
if items == nil {
items = []ScheduleTaskResponse{}
}
return &ScheduleTaskListResponse{Items: items, Total: total}, nil
}
func (s *ScheduleTaskService) Detail(ctx context.Context, id int64) (*ScheduleTaskResponse, error) {
actor := auth.MustActor(ctx)
var r ScheduleTaskResponse
var ca, ua interface{}
var startAt, endAt, nextRunAt interface{}
err := s.pool.QueryRow(ctx, `
SELECT st.id, st.prompt_rule_id, st.brand_id, st.name,
st.cron_expr, st.target_platform, st.enable_web_search, st.generate_count, st.start_at, st.end_at,
st.next_run_at, st.status, st.created_at, st.updated_at,
pr.name AS prompt_rule_name,
b.name AS brand_name
FROM schedule_tasks st
LEFT JOIN prompt_rules pr ON pr.id = st.prompt_rule_id
LEFT JOIN brands b ON b.id = st.brand_id
WHERE st.id = $1 AND st.tenant_id = $2 AND st.deleted_at IS NULL
`, id, actor.TenantID).Scan(&r.ID, &r.PromptRuleID, &r.BrandID, &r.Name,
&r.CronExpr, &r.TargetPlatform, &r.EnableWebSearch, &r.GenerateCount, &startAt, &endAt,
&nextRunAt, &r.Status, &ca, &ua,
&r.PromptRuleName, &r.BrandName)
if err != nil {
return nil, response.ErrNotFound(40432, "schedule_not_found", "schedule task not found")
}
r.CreatedAt = fmt.Sprintf("%v", ca)
r.UpdatedAt = fmt.Sprintf("%v", ua)
if startAt != nil {
sv := fmt.Sprintf("%v", startAt)
r.StartAt = &sv
}
if endAt != nil {
sv := fmt.Sprintf("%v", endAt)
r.EndAt = &sv
}
if nextRunAt != nil {
sv := fmt.Sprintf("%v", nextRunAt)
r.NextRunAt = &sv
}
return &r, nil
}
func (s *ScheduleTaskService) Create(ctx context.Context, req ScheduleTaskRequest) (*ScheduleTaskResponse, error) {
actor := auth.MustActor(ctx)
if err := validateScheduleCronExpr(req.CronExpr); err != nil {
return nil, err
}
if err := s.ensurePromptRuleExists(ctx, actor.TenantID, req.PromptRuleID); err != nil {
return nil, err
}
generateCount, err := normalizeScheduleGenerateCount(req.GenerateCount)
if err != nil {
return nil, err
}
enableWebSearch := boolValue(req.EnableWebSearch)
tx, err := s.pool.Begin(ctx)
if err != nil {
return nil, response.ErrInternal(50010, "create_failed", "failed to begin schedule task transaction")
}
defer tx.Rollback(ctx)
var id int64
var createdAt time.Time
var startAt pgtype.Timestamptz
var endAt pgtype.Timestamptz
var status string
err = tx.QueryRow(ctx, `
INSERT INTO schedule_tasks (
tenant_id, operator_id, prompt_rule_id, brand_id, name, cron_expr, target_platform,
enable_web_search, generate_count, start_at, end_at
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::timestamptz, $11::timestamptz)
RETURNING id, created_at, start_at, end_at, status
`, actor.TenantID, actor.UserID, req.PromptRuleID, req.BrandID, req.Name, req.CronExpr, req.TargetPlatform, enableWebSearch, generateCount, req.StartAt, req.EndAt).
Scan(&id, &createdAt, &startAt, &endAt, &status)
if err != nil {
return nil, response.ErrInternal(50010, "create_failed", "failed to create schedule task")
}
nextRunAt, err := resolveScheduleNextRunForStatus(status, req.CronExpr, timePtrFromTimestamp(startAt), timePtrFromTimestamp(endAt), time.Now())
if err != nil {
return nil, response.ErrBadRequest(40015, "invalid_cron_expr", err.Error())
}
if _, err := tx.Exec(ctx, `
UPDATE schedule_tasks
SET next_run_at = $1,
updated_at = NOW()
WHERE id = $2
AND tenant_id = $3
`, nextRunAt, id, actor.TenantID); err != nil {
return nil, response.ErrInternal(50010, "create_failed", "failed to initialize next run time")
}
if err := tx.Commit(ctx); err != nil {
return nil, response.ErrInternal(50010, "create_failed", "failed to commit schedule task")
}
afterJSON, _ := json.Marshal(map[string]interface{}{
"id": id,
"name": req.Name,
"cron_expr": req.CronExpr,
"enable_web_search": enableWebSearch,
"generate_count": generateCount,
})
result := "success"
resourceType := "schedule_task"
requestID := middleware.RequestIDFromContext(ctx)
s.auditLogs.Log(auditlog.Entry{
OperatorID: actor.UserID,
TenantID: &actor.TenantID,
Module: "schedule_task",
Action: "create",
ResourceType: &resourceType,
ResourceID: &id,
RequestID: nilIfEmptyString(requestID),
AfterJSON: afterJSON,
Result: &result,
})
return &ScheduleTaskResponse{
ID: id, PromptRuleID: req.PromptRuleID, BrandID: req.BrandID,
Name: req.Name, CronExpr: req.CronExpr, TargetPlatform: req.TargetPlatform,
EnableWebSearch: enableWebSearch, GenerateCount: generateCount,
Status: status, CreatedAt: fmt.Sprintf("%v", createdAt),
NextRunAt: timeStringPtr(nextRunAt),
}, nil
}
func (s *ScheduleTaskService) Update(ctx context.Context, id int64, req ScheduleTaskRequest) error {
actor := auth.MustActor(ctx)
if err := validateScheduleCronExpr(req.CronExpr); err != nil {
return err
}
if err := s.ensurePromptRuleExists(ctx, actor.TenantID, req.PromptRuleID); err != nil {
return err
}
generateCount, err := normalizeScheduleGenerateCount(req.GenerateCount)
if err != nil {
return err
}
enableWebSearch := boolValue(req.EnableWebSearch)
tx, err := s.pool.Begin(ctx)
if err != nil {
return response.ErrInternal(50010, "update_failed", "failed to begin schedule task transaction")
}
defer tx.Rollback(ctx)
var status string
var startAt pgtype.Timestamptz
var endAt pgtype.Timestamptz
err = tx.QueryRow(ctx, `
UPDATE schedule_tasks SET prompt_rule_id = $1, brand_id = $2, name = $3,
cron_expr = $4, target_platform = $5, enable_web_search = $6, generate_count = $7,
start_at = $8::timestamptz, end_at = $9::timestamptz, operator_id = $10, updated_at = NOW()
WHERE id = $11 AND tenant_id = $12 AND deleted_at IS NULL
RETURNING status, start_at, end_at
`, req.PromptRuleID, req.BrandID, req.Name, req.CronExpr, req.TargetPlatform, enableWebSearch, generateCount, req.StartAt, req.EndAt, actor.UserID, id, actor.TenantID).
Scan(&status, &startAt, &endAt)
if err != nil {
return response.ErrNotFound(40432, "schedule_not_found", "schedule task not found")
}
nextRunAt, err := resolveScheduleNextRunForStatus(status, req.CronExpr, timePtrFromTimestamp(startAt), timePtrFromTimestamp(endAt), time.Now())
if err != nil {
return response.ErrBadRequest(40015, "invalid_cron_expr", err.Error())
}
if _, err := tx.Exec(ctx, `
UPDATE schedule_tasks
SET next_run_at = $1,
updated_at = NOW()
WHERE id = $2
AND tenant_id = $3
`, nextRunAt, id, actor.TenantID); err != nil {
return response.ErrInternal(50010, "update_failed", "failed to refresh next run time")
}
if err := tx.Commit(ctx); err != nil {
return response.ErrInternal(50010, "update_failed", "failed to commit schedule task")
}
return nil
}
func (s *ScheduleTaskService) Delete(ctx context.Context, id int64) error {
actor := auth.MustActor(ctx)
tag, err := s.pool.Exec(ctx, `
UPDATE schedule_tasks SET deleted_at = NOW(), updated_at = NOW()
WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL
`, id, actor.TenantID)
if err != nil || tag.RowsAffected() == 0 {
return response.ErrNotFound(40432, "schedule_not_found", "schedule task not found")
}
afterJSON, _ := json.Marshal(map[string]interface{}{"id": id})
result := "success"
resourceType := "schedule_task"
requestID := middleware.RequestIDFromContext(ctx)
s.auditLogs.Log(auditlog.Entry{
OperatorID: actor.UserID,
TenantID: &actor.TenantID,
Module: "schedule_task",
Action: "delete",
ResourceType: &resourceType,
ResourceID: &id,
RequestID: nilIfEmptyString(requestID),
AfterJSON: afterJSON,
Result: &result,
})
return nil
}
func (s *ScheduleTaskService) ToggleStatus(ctx context.Context, id int64, req ScheduleTaskStatusRequest) error {
actor := auth.MustActor(ctx)
tx, err := s.pool.Begin(ctx)
if err != nil {
return response.ErrInternal(50010, "update_failed", "failed to begin schedule task transaction")
}
defer tx.Rollback(ctx)
var cronExpr string
var status string
var startAt pgtype.Timestamptz
var endAt pgtype.Timestamptz
err = tx.QueryRow(ctx, `
UPDATE schedule_tasks SET status = $1, operator_id = $2, updated_at = NOW()
WHERE id = $3 AND tenant_id = $4 AND deleted_at IS NULL
RETURNING cron_expr, status, start_at, end_at
`, req.Status, actor.UserID, id, actor.TenantID).
Scan(&cronExpr, &status, &startAt, &endAt)
if err != nil {
return response.ErrNotFound(40432, "schedule_not_found", "schedule task not found")
}
nextRunAt, err := resolveScheduleNextRunForStatus(status, cronExpr, timePtrFromTimestamp(startAt), timePtrFromTimestamp(endAt), time.Now())
if err != nil {
return response.ErrBadRequest(40015, "invalid_cron_expr", err.Error())
}
if _, err := tx.Exec(ctx, `
UPDATE schedule_tasks
SET next_run_at = $1,
updated_at = NOW()
WHERE id = $2
AND tenant_id = $3
`, nextRunAt, id, actor.TenantID); err != nil {
return response.ErrInternal(50010, "update_failed", "failed to refresh next run time")
}
if err := tx.Commit(ctx); err != nil {
return response.ErrInternal(50010, "update_failed", "failed to commit schedule task")
}
return nil
}
func (s *ScheduleTaskService) ensurePromptRuleExists(ctx context.Context, tenantID, promptRuleID int64) error {
var exists bool
_ = s.pool.QueryRow(ctx, `
SELECT EXISTS(SELECT 1 FROM prompt_rules WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL)
`, promptRuleID, tenantID).Scan(&exists)
if !exists {
return response.ErrBadRequest(40001, "invalid_rule", "prompt rule not found or not active")
}
return nil
}
func resolveScheduleNextRunForStatus(status, cronExpr string, startAt, endAt *time.Time, now time.Time) (*time.Time, error) {
if status != "enabled" {
return nil, nil
}
return sharedschedule.ComputeNextRun(cronExpr, startAt, endAt, now)
}
func validateScheduleCronExpr(cronExpr string) error {
if _, err := sharedschedule.Parse(cronExpr); err != nil {
return response.ErrBadRequest(40015, "invalid_cron_expr", err.Error())
}
return nil
}
func normalizeScheduleGenerateCount(value *int) (int, error) {
if value == nil {
return 1, nil
}
if *value < 1 || *value > 20 {
return 0, response.ErrBadRequest(40022, "invalid_generate_count", "generate_count must be between 1 and 20")
}
return *value, nil
}
func boolValue(value *bool) bool {
return value != nil && *value
}
func timePtrFromTimestamp(value pgtype.Timestamptz) *time.Time {
if !value.Valid {
return nil
}
result := value.Time.Round(0)
return &result
}
func timeStringPtr(value *time.Time) *string {
if value == nil {
return nil
}
text := value.Format(time.RFC3339)
return &text
}