orbit/internal/storage/mysql/task_query.go

package mysql

import (
	"context"
	"database/sql"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"time"

	"git.yuntuo.hk/k2safe/orbit/internal/fault"
	"git.yuntuo.hk/k2safe/orbit/internal/jsonvalue"
	core "git.yuntuo.hk/k2safe/orbit/model"
	"git.yuntuo.hk/k2safe/orbit/platform/auth"
	pulse "git.yuntuo.hk/k2safe/orbit/pulse/model"
	"git.yuntuo.hk/k2safe/orbit/pulse/service"
)

// taskCursor 是已授权查询的排序位置，不能代替凭据；v2摘要绑定全部筛选。
type taskCursor struct {
	Version   int             `json:"v,omitempty"`
	Namespace core.ID         `json:"n"`
	Queue     core.ID         `json:"q"`
	State     pulse.TaskState `json:"s"`
	Created   time.Time       `json:"t"`
	ID        core.ID         `json:"i"`
	Filter    string          `json:"f,omitempty"`
}

func taskFilterDigest(in service.ListInput) (string, error) {
	digest, err := requestDigest(struct {
		Domain      string          `json:"domain"`
		Queue       core.ID         `json:"queue"`
		State       pulse.TaskState `json:"state"`
		BusinessKey string          `json:"business_key"`
		From        *time.Time      `json:"from"`
		Before      *time.Time      `json:"before"`
	}{"orbit.task-query.v2", in.QueueID, in.State, in.BusinessKey, in.CreatedFrom, in.CreatedBefore})
	return hex.EncodeToString(digest[:]), err
}

func decodeTaskCursor(value string, ns core.ID, in service.ListInput, filter string) (taskCursor, error) {
	var c taskCursor
	invalid := fault.Invalid("cursor does not match this task query")
	raw, err := base64.RawURLEncoding.Strict().DecodeString(value)
	if err != nil || len(raw) > 768 {
		return c, invalid
	}
	if _, err = jsonvalue.Canonical(raw); err != nil {
		return c, invalid
	}
	var fields map[string]json.RawMessage
	if json.Unmarshal(raw, &fields) != nil || fields == nil {
		return c, invalid
	}
	for k, v := range fields {
		switch k {
		case "v", "n", "q", "s", "t", "i", "f":
		default:
			return c, invalid
		}
		if string(v) == "null" {
			return c, invalid
		}
	}
	for _, k := range []string{"n", "q", "s", "t", "i"} {
		if _, ok := fields[k]; !ok {
			return c, invalid
		}
	}
	if json.Unmarshal(raw, &c) != nil || c.Namespace != ns || c.Queue != in.QueueID || c.State != in.State || c.ID.Validate() != nil || c.Created.Year() < 2000 || c.Created.Year() > 9998 || c.Created.Nanosecond()%1000 != 0 {
		return c, invalid
	}
	if c.Version == 0 {
		// 旧版未带版本的游标仅兼容原队列/状态查询，不能附加新筛选。
		if _, ok := fields["v"]; ok {
			return c, invalid
		}
		if _, ok := fields["f"]; ok {
			return c, invalid
		}
		if in.BusinessKey != "" || in.CreatedFrom != nil || in.CreatedBefore != nil {
			return c, invalid
		}
	} else if c.Version != 2 || c.Filter != filter {
		return c, invalid
	}
	if in.CreatedFrom != nil && c.Created.Before(*in.CreatedFrom) || in.CreatedBefore != nil && !c.Created.Before(*in.CreatedBefore) {
		return c, invalid
	}
	return c, nil
}

const taskSummaryColumns = `id,queue_id,state,revision,handler,handler_version,COALESCE(business_key,''),COALESCE(concurrency_key,''),priority,available_at,deadline_at,attempt_count,created_at,updated_at,effect_uncertain`

// taskListSQL 选择固定索引名，等值筛选全部位于创建时间前；游标不依赖状态更新时间。
func taskListSQL(ns core.ID, in service.ListInput, filter string) (string, []any, error) {
	index := "ix_task_created"
	if in.State != "" {
		index = "ix_task_state_created"
	}
	if in.BusinessKey != "" {
		index = "ix_task_queue_business_created"
		if in.State != "" {
			index = "ix_task_queue_business_state_created"
		}
	}
	q := `SELECT ` + taskSummaryColumns + ` FROM pulse_tasks FORCE INDEX (` + index + `) WHERE namespace_id=? AND queue_id=?`
	args := []any{ns, in.QueueID}
	if in.BusinessKey != "" {
		q += " AND business_key=?"
		args = append(args, in.BusinessKey)
	}
	if in.State != "" {
		q += " AND state=?"
		args = append(args, in.State)
	}
	if in.CreatedFrom != nil {
		q += " AND created_at>=?"
		args = append(args, *in.CreatedFrom)
	}
	if in.CreatedBefore != nil {
		q += " AND created_at<?"
		args = append(args, *in.CreatedBefore)
	}
	if in.Cursor != "" {
		c, err := decodeTaskCursor(in.Cursor, ns, in, filter)
		if err != nil {
			return "", nil, err
		}
		q += " AND (created_at<? OR (created_at=? AND id<?))"
		args = append(args, c.Created, c.Created, c.ID)
	}
	q += " ORDER BY created_at DESC,id DESC LIMIT ?"
	args = append(args, in.Limit+1)
	return q, args, nil
}

func (s *TaskStore) ListTasks(ctx context.Context, token auth.Token, in service.ListInput) (service.TaskPage, error) {
	out := service.TaskPage{Items: []service.TaskSummary{}}
	err := s.transaction(ctx, token, "tasks:read", func(tx *sql.Tx, p auth.Permissions) error {
		if !p.AllowsQueue(in.QueueID) {
			return fault.Forbidden
		}
		filter, err := taskFilterDigest(in)
		if err != nil {
			return err
		}
		q, args, err := taskListSQL(token.NamespaceID, in, filter)
		if err != nil {
			return err
		}
		rows, err := tx.QueryContext(ctx, q, args...)
		if err != nil {
			return err
		}
		defer rows.Close()
		for rows.Next() {
			var v service.TaskSummary
			if err := rows.Scan(&v.ID, &v.QueueID, &v.State, &v.Revision, &v.Handler, &v.HandlerVersion, &v.BusinessKey, &v.ConcurrencyKey, &v.Priority, &v.AvailableAt, &v.DeadlineAt, &v.AttemptCount, &v.CreatedAt, &v.UpdatedAt, &v.EffectUncertain); err != nil {
				return err
			}
			out.Items = append(out.Items, v)
		}
		if err := rows.Err(); err != nil {
			return err
		}
		if len(out.Items) > in.Limit {
			out.Items = out.Items[:in.Limit]
			last := out.Items[len(out.Items)-1]
			raw, err := json.Marshal(taskCursor{Version: 2, Namespace: token.NamespaceID, Queue: in.QueueID, State: in.State, Created: last.CreatedAt, ID: last.ID, Filter: filter})
			if err != nil {
				return err
			}
			out.NextCursor = base64.RawURLEncoding.EncodeToString(raw)
		}
		return nil
	})
	return out, err
}