goext/cursortoken/token.go

package cursortoken

import (
	"encoding/base32"
	"encoding/json"
	"errors"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"strings"
	"time"
)

type Mode string

const (
	CTMStart  Mode = "START"
	CTMNormal Mode = "NORMAL"
	CTMEnd    Mode = "END"
)

type Extra struct {
	Timestamp *time.Time
	Id        *string
	Page      *int
	PageSize  *int
}

type CursorToken struct {
	Mode               Mode
	ValuePrimary       string
	ValueSecondary     string
	Direction          SortDirection
	DirectionSecondary SortDirection
	PageSize           int
	Extra              Extra
}

type cursorTokenSerialize struct {
	ValuePrimary       *string        `json:"v1,omitempty"`
	ValueSecondary     *string        `json:"v2,omitempty"`
	Direction          *SortDirection `json:"dir,omitempty"`
	DirectionSecondary *SortDirection `json:"dir2,omitempty"`
	PageSize           *int           `json:"size,omitempty"`

	ExtraTimestamp *time.Time `json:"ts,omitempty"`
	ExtraId        *string    `json:"id,omitempty"`
	ExtraPage      *int       `json:"pg,omitempty"`
	ExtraPageSize  *int       `json:"sz,omitempty"`
}

func Start() CursorToken {
	return CursorToken{
		Mode:               CTMStart,
		ValuePrimary:       "",
		ValueSecondary:     "",
		Direction:          "",
		DirectionSecondary: "",
		PageSize:           0,
		Extra:              Extra{},
	}
}

func End() CursorToken {
	return CursorToken{
		Mode:               CTMEnd,
		ValuePrimary:       "",
		ValueSecondary:     "",
		Direction:          "",
		DirectionSecondary: "",
		PageSize:           0,
		Extra:              Extra{},
	}
}

func (c *CursorToken) Token() string {
	if c.Mode == CTMStart {
		return "@start"
	}
	if c.Mode == CTMEnd {
		return "@end"
	}

	// We kinda manually implement omitempty for the CursorToken here
	// because omitempty does not work for time.Time and otherwise we would always
	// get weird time values when decoding a token that initially didn't have an Timestamp set
	// For this usecase we treat Unix=0 as an empty timestamp

	sertok := cursorTokenSerialize{}

	if c.ValuePrimary != "" {
		sertok.ValuePrimary = &c.ValuePrimary
	}
	if c.ValueSecondary != "" {
		sertok.ValueSecondary = &c.ValueSecondary
	}
	if c.Direction != "" {
		sertok.Direction = &c.Direction
	}
	if c.DirectionSecondary != "" {
		sertok.DirectionSecondary = &c.DirectionSecondary
	}
	if c.PageSize != 0 {
		sertok.PageSize = &c.PageSize
	}

	sertok.ExtraTimestamp = c.Extra.Timestamp
	sertok.ExtraId = c.Extra.Id
	sertok.ExtraPage = c.Extra.Page
	sertok.ExtraPageSize = c.Extra.PageSize

	body, err := json.Marshal(sertok)
	if err != nil {
		panic(err)
	}

	return "tok_" + base32.StdEncoding.EncodeToString(body)
}

func Decode(tok string) (CursorToken, error) {
	if tok == "" {
		return Start(), nil
	}
	if strings.ToLower(tok) == "@start" {
		return Start(), nil
	}
	if strings.ToLower(tok) == "@end" {
		return End(), nil
	}

	if !strings.HasPrefix(tok, "tok_") {
		return CursorToken{}, errors.New("could not decode token, missing prefix")
	}

	body, err := base32.StdEncoding.DecodeString(tok[len("tok_"):])
	if err != nil {
		return CursorToken{}, err
	}

	var tokenDeserialize cursorTokenSerialize
	err = json.Unmarshal(body, &tokenDeserialize)
	if err != nil {
		return CursorToken{}, err
	}

	token := CursorToken{Mode: CTMNormal}

	if tokenDeserialize.ValuePrimary != nil {
		token.ValuePrimary = *tokenDeserialize.ValuePrimary
	}
	if tokenDeserialize.ValueSecondary != nil {
		token.ValueSecondary = *tokenDeserialize.ValueSecondary
	}
	if tokenDeserialize.Direction != nil {
		token.Direction = *tokenDeserialize.Direction
	}
	if tokenDeserialize.DirectionSecondary != nil {
		token.DirectionSecondary = *tokenDeserialize.DirectionSecondary
	}
	if tokenDeserialize.PageSize != nil {
		token.PageSize = *tokenDeserialize.PageSize
	}

	token.Extra.Timestamp = tokenDeserialize.ExtraTimestamp
	token.Extra.Id = tokenDeserialize.ExtraId
	token.Extra.Page = tokenDeserialize.ExtraPage
	token.Extra.PageSize = tokenDeserialize.ExtraPageSize

	return token, nil
}

func (c *CursorToken) ValuePrimaryObjectId() (primitive.ObjectID, bool) {
	if oid, err := primitive.ObjectIDFromHex(c.ValuePrimary); err == nil {
		return oid, true
	} else {
		return primitive.ObjectID{}, false
	}
}

func (c *CursorToken) ValueSecondaryObjectId() (primitive.ObjectID, bool) {
	if oid, err := primitive.ObjectIDFromHex(c.ValueSecondary); err == nil {
		return oid, true
	} else {
		return primitive.ObjectID{}, false
	}
}