package queryspec

import "context"

type aggregateFetchLimitKey struct{}

// WithAggregateFetchLimit asks aggregate query execution to fetch enough grouped
// rows for a caller-owned result window before the caller slices a physical page
// out of that window, or to fetch a one-row sentinel when the caller needs to
// detect that a grouped result would exceed one REST batch.
func WithAggregateFetchLimit(ctx context.Context, limit int) context.Context {
	if limit <= 0 {
		return ctx
	}
	return context.WithValue(ctx, aggregateFetchLimitKey{}, limit)
}

// AggregateFetchLimitFromContext returns the aggregate fetch limit requested by
// WithAggregateFetchLimit.
func AggregateFetchLimitFromContext(ctx context.Context) int {
	limit, _ := ctx.Value(aggregateFetchLimitKey{}).(int)
	return limit
}