cursor/query.go

package cursor

import (
	"fmt"
	"log/slog"

	"gitea.auvem.com/go-toolkit/dbx"
	"github.com/go-jet/jet/v2/mysql"
)

// QueryCountResult is a struct that contains the total number of rows that will
// be returned by a query, the number of rows after an end row, and the number of
// rows before a start row.
type QueryCountResult struct {
	Total  int
	Before int
	After  int
}

// QueryCountFn is a function type that abstracts the counting of rows down to
// a single function that takes a Queryable interface and returns a QueryCountResult.
type QueryCountFn = func(
	sqlo dbx.Queryable,
	start GenericCursor,
	end GenericCursor,
) (QueryCountResult, error)

// QueryCount counts the number of rows returned from a table with the given
// conditions. Returns a QueryCountResult or an error if anything goes wrong.
func QueryCount(
	sqlo dbx.Queryable,
	col mysql.Column,
	tbl mysql.ReadableTable,
	conds mysql.BoolExpression,
	start GenericCursor,
	end GenericCursor,
) (QueryCountResult, error) {
	var comparator struct {
		start string
		end   string
	}

	if start.Direction() == OrderAscending {
		comparator.start = "<"
		comparator.end = ">"
	} else {
		comparator.start = ">"
		comparator.end = "<"
	}

	stmt := tbl.SELECT(
		mysql.COUNT(col).AS("QueryCountResult.Total"),
		mysql.COUNT(countBoundExpr(start, comparator.start)).AS("QueryCountResult.Before"),
		mysql.COUNT(countBoundExpr(end, comparator.end)).AS("QueryCountResult.After"),
	).WHERE(conds)
	var res QueryCountResult
	err := stmt.Query(sqlo, &res)
	if err != nil {
		return QueryCountResult{}, fmt.Errorf("failed to query count: %w", err)
	}
	return res, nil
}

func countBoundExpr(c GenericCursor, comparator string) mysql.Expression {
	if c.IsComposite() {
		return mysql.Raw(
			fmt.Sprintf("IF(%s, 1, NULL)", compositeTupleSQL(comparator, c)),
			compositeTupleArgs(c),
		)
	}

	ck := c.GenericIndex().ColumnKey()
	return mysql.Raw(
		fmt.Sprintf("IF(%s.%s %s ?, 1, NULL)", ck.Table, ck.Column, comparator),
		mysql.RawArgs{"?": c.GenericIndex()},
	)
}

func compositeTupleSQL(comparator string, c GenericCursor) string {
	orderCK := NewColumnKey(c.OrderCol())
	indexCK := c.GenericIndex().ColumnKey()
	return fmt.Sprintf(
		"(%s.%s, %s.%s) %s (#order, #index)",
		orderCK.Table, orderCK.Column,
		indexCK.Table, indexCK.Column,
		comparator,
	)
}

func compositeTupleArgs(c GenericCursor) mysql.RawArgs {
	return mysql.RawArgs{
		"#order": c.GenericOrderValue(),
		"#index": c.GenericIndex(),
	}
}

// BuildQueryCountFn builds a QueryCountFn that can be used to count rows in a
// table with the given conditions.
func BuildQueryCountFn(
	col mysql.Column,
	tbl mysql.ReadableTable,
	conds mysql.BoolExpression,
) QueryCountFn {
	return func(sqlo dbx.Queryable, start GenericCursor, end GenericCursor) (QueryCountResult, error) {
		return QueryCount(sqlo, col, tbl, conds, start, end)
	}
}

// PaginateConds returns a mysql.BoolExpression that paginates results using the
// provided cursor as a base position. Nil or empty cursors match all rows.
func PaginateConds[IE mysql.Expression, IC mysql.Column](c *Cursor[IE, IC]) mysql.BoolExpression {
	return paginateFromGeneric(c)
}

func paginateFromGeneric(c GenericCursor) mysql.BoolExpression {
	if c == nil || c.IsEmpty() {
		return mysql.Bool(true)
	}

	if c.IsComposite() {
		return paginateComposite(c)
	}

	switch idx := c.GenericIndex().(type) {
	case *Int64Value:
		if c.Direction() == OrderAscending {
			return idx.Col().GT(idx.Expr())
		}
		return idx.Col().LT(idx.Expr())
	case *Uint64Value:
		if c.Direction() == OrderAscending {
			return idx.Col().GT(idx.Expr())
		}
		return idx.Col().LT(idx.Expr())
	case *StringValue:
		if c.Direction() == OrderAscending {
			return idx.Col().GT(idx.Expr())
		}
		return idx.Col().LT(idx.Expr())
	case *TimestampValue:
		if c.Direction() == OrderAscending {
			return idx.Col().GT(idx.Expr())
		}
		return idx.Col().LT(idx.Expr())
	default:
		// fallback: emit a raw comparison using the column key and the value
		ck := c.GenericIndex().ColumnKey()
		if ck.IsEmpty() {
			return mysql.Bool(true)
		}
		slog.Warn(
			"Cursor pagination conditions generated via raw fallback",
			"table", ck.Table, "col", ck.Column,
		)
		if c.Direction() == OrderAscending {
			return mysql.RawBool(
				fmt.Sprintf("%s.%s > ?", ck.Table, ck.Column),
				mysql.RawArgs{"?": c.GenericIndex()},
			)
		}
		return mysql.RawBool(
			fmt.Sprintf("%s.%s < ?", ck.Table, ck.Column),
			mysql.RawArgs{"?": c.GenericIndex()},
		)
	}
}

func paginateComposite(c GenericCursor) mysql.BoolExpression {
	comparator := "<"
	if c.Direction() == OrderAscending {
		comparator = ">"
	}

	return mysql.RawBool(
		compositeTupleSQL(comparator, c),
		compositeTupleArgs(c),
	)
}

// OrderByClauses returns ORDER BY clauses matching the cursor pagination semantics,
// including tuple ordering when [GenericCursor.UsesTupleOrdering] is true.
func OrderByClauses(c GenericCursor) []mysql.OrderByClause {
	if c == nil {
		return nil
	}

	orderCol := c.OrderCol()
	if c.UsesTupleOrdering() {
		indexCol, err := GetColumnByKey(c.GenericIndex().ColumnKey())
		if err != nil {
			panic(err)
		}
		indexOrderCol, ok := indexCol.(CursorOrderCol)
		if !ok {
			panic(fmt.Errorf("index column %s is not orderable", c.GenericIndex().ColumnKey()))
		}
		if c.Direction() == OrderAscending {
			return []mysql.OrderByClause{orderCol.ASC(), indexOrderCol.ASC()}
		}
		return []mysql.OrderByClause{orderCol.DESC(), indexOrderCol.DESC()}
	}

	if c.Direction() == OrderAscending {
		return []mysql.OrderByClause{orderCol.ASC()}
	}
	return []mysql.OrderByClause{orderCol.DESC()}
}