fathom/internal/query/parse.go

// Package query implements Fathom's deliberately small SQL subset.
//
// The grammar is:
//
//	SELECT col, agg(col), ...
//	FROM table
//	WHERE time >= min AND time <= max [AND cond ...]
//	[GROUP BY col, ...]
//	[ORDER BY expr [ASC|DESC]]
//	[LIMIT n]
//
// No joins, no subqueries, no window functions, no OR. A trace is
// self-contained — parent_span_id rebuilds the tree within one table — so there
// is nothing to join against, and the analysis workload does not need the rest.
// Keeping the grammar this small is what lets the executor be a scan and a
// hash-aggregate with no planner worth the name.
package query

import (
	"fmt"
	"strconv"
	"strings"
)

// AggFunc is a supported aggregate.
type AggFunc string

const (
	AggNone  AggFunc = ""
	AggCount AggFunc = "count"
	AggSum   AggFunc = "sum"
	AggAvg   AggFunc = "avg"
	AggMin   AggFunc = "min"
	AggMax   AggFunc = "max"
	AggP50   AggFunc = "p50"
	AggP90   AggFunc = "p90"
	AggP95   AggFunc = "p95"
	AggP99   AggFunc = "p99"
)

func (a AggFunc) isPercentile() bool {
	switch a {
	case AggP50, AggP90, AggP95, AggP99:
		return true
	}
	return false
}

func (a AggFunc) percentile() float64 {
	switch a {
	case AggP50:
		return 0.50
	case AggP90:
		return 0.90
	case AggP95:
		return 0.95
	case AggP99:
		return 0.99
	}
	return 0
}

// Select is one item in the projection.
type Select struct {
	Agg    AggFunc
	Column string // "*" only with count
	Alias  string
}

// Label is the column header this item produces.
func (s Select) Label() string {
	if s.Alias != "" {
		return s.Alias
	}
	if s.Agg == AggNone {
		return s.Column
	}
	return string(s.Agg) + "(" + s.Column + ")"
}

// Op is a comparison operator.
type Op string

const (
	OpEq       Op = "="
	OpNe       Op = "!="
	OpLt       Op = "<"
	OpLe       Op = "<="
	OpGt       Op = ">"
	OpGe       Op = ">="
	OpLike     Op = "like"
	OpContains Op = "contains"
)

// Condition is a single WHERE predicate. Conditions are ANDed; there is no OR,
// which keeps the scan a straight-line filter.
type Condition struct {
	Column string
	Op     Op
	Value  string
	Num    float64
	IsNum  bool
}

// Query is a parsed statement.
type Query struct {
	Selects []Select
	Table   string
	Conds   []Condition
	MinTS   int64 // from the mandatory time range, unix micros
	MaxTS   int64
	GroupBy []string
	OrderBy string
	Desc    bool
	Limit   int
}

// HasAggregates reports whether this is a grouped/aggregated query rather than
// a row listing.
func (q *Query) HasAggregates() bool {
	for _, s := range q.Selects {
		if s.Agg != AggNone {
			return true
		}
	}
	return false
}

// Parse compiles a statement.
func Parse(src string) (*Query, error) {
	toks, err := lex(src)
	if err != nil {
		return nil, err
	}
	p := &parser{toks: toks}
	q, err := p.parse()
	if err != nil {
		return nil, err
	}
	return q, nil
}

type tokKind int

const (
	tokIdent tokKind = iota
	tokNumber
	tokString
	tokPunct
)

type token struct {
	kind tokKind
	s    string
	pos  int
}

func lex(src string) ([]token, error) {
	var out []token
	i := 0
	for i < len(src) {
		c := src[i]
		switch {
		case c == ' ' || c == '\t' || c == '\n' || c == '\r':
			i++
		case c == ',' || c == '(' || c == ')' || c == '*':
			out = append(out, token{tokPunct, string(c), i})
			i++
		case c == '>' || c == '<' || c == '!' || c == '=':
			if i+1 < len(src) && src[i+1] == '=' {
				out = append(out, token{tokPunct, src[i : i+2], i})
				i += 2
			} else {
				out = append(out, token{tokPunct, string(c), i})
				i++
			}
		case c == '\'' || c == '"':
			quote := c
			j := i + 1
			var sb strings.Builder
			for j < len(src) && src[j] != quote {
				// Backslash escapes, so a literal can contain its own quote.
				if src[j] == '\\' && j+1 < len(src) {
					sb.WriteByte(src[j+1])
					j += 2
					continue
				}
				sb.WriteByte(src[j])
				j++
			}
			if j >= len(src) {
				return nil, fmt.Errorf("unterminated string starting at position %d", i)
			}
			out = append(out, token{tokString, sb.String(), i})
			i = j + 1
		case c >= '0' && c <= '9' || c == '-' && i+1 < len(src) && src[i+1] >= '0' && src[i+1] <= '9':
			j := i + 1
			for j < len(src) && (src[j] >= '0' && src[j] <= '9' || src[j] == '.') {
				j++
			}
			out = append(out, token{tokNumber, src[i:j], i})
			i = j
		case isIdentChar(c):
			j := i
			for j < len(src) && isIdentChar(src[j]) {
				j++
			}
			out = append(out, token{tokIdent, src[i:j], i})
			i = j
		default:
			return nil, fmt.Errorf("unexpected character %q at position %d", c, i)
		}
	}
	return out, nil
}

// isIdentChar accepts dots so that attribute names like gen_ai.request.model
// are single identifiers rather than needing quoting.
func isIdentChar(c byte) bool {
	return c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' ||
		c >= '0' && c <= '9' || c == '_' || c == '.'
}

type parser struct {
	toks []token
	i    int
}

func (p *parser) peek() (token, bool) {
	if p.i < len(p.toks) {
		return p.toks[p.i], true
	}
	return token{}, false
}

func (p *parser) next() (token, bool) {
	t, ok := p.peek()
	if ok {
		p.i++
	}
	return t, ok
}

// acceptKeyword consumes a case-insensitive keyword if present.
func (p *parser) acceptKeyword(kw string) bool {
	t, ok := p.peek()
	if ok && t.kind == tokIdent && strings.EqualFold(t.s, kw) {
		p.i++
		return true
	}
	return false
}

func (p *parser) expectKeyword(kw string) error {
	if p.acceptKeyword(kw) {
		return nil
	}
	if t, ok := p.peek(); ok {
		return fmt.Errorf("expected %s at position %d, found %q", strings.ToUpper(kw), t.pos, t.s)
	}
	return fmt.Errorf("expected %s but the statement ended", strings.ToUpper(kw))
}

func (p *parser) acceptPunct(s string) bool {
	t, ok := p.peek()
	if ok && t.kind == tokPunct && t.s == s {
		p.i++
		return true
	}
	return false
}

func (p *parser) parse() (*Query, error) {
	q := &Query{Limit: 100}

	if err := p.expectKeyword("select"); err != nil {
		return nil, err
	}
	sels, err := p.parseSelects()
	if err != nil {
		return nil, err
	}
	q.Selects = sels

	if err := p.expectKeyword("from"); err != nil {
		return nil, err
	}
	t, ok := p.next()
	if !ok || t.kind != tokIdent {
		return nil, fmt.Errorf("expected a table name after FROM")
	}
	q.Table = t.s

	// WHERE is mandatory, because the time range is mandatory. Without it a
	// query would scan every shard that exists, which is the one thing a store
	// with no index other than time cannot afford to make easy.
	if err := p.expectKeyword("where"); err != nil {
		return nil, fmt.Errorf("%w (a time range is required: WHERE time >= ... AND time <= ...)", err)
	}
	if err := p.parseWhere(q); err != nil {
		return nil, err
	}
	if q.MinTS == 0 && q.MaxTS == 0 {
		return nil, fmt.Errorf("the WHERE clause must bound `time`: every query needs a time range")
	}
	if q.MaxTS == 0 {
		q.MaxTS = 1<<62 - 1
	}
	if q.MinTS > q.MaxTS {
		return nil, fmt.Errorf("the time range is inverted: min %d is after max %d", q.MinTS, q.MaxTS)
	}

	if p.acceptKeyword("group") {
		if err := p.expectKeyword("by"); err != nil {
			return nil, err
		}
		for {
			t, ok := p.next()
			if !ok || t.kind != tokIdent {
				return nil, fmt.Errorf("expected a column name in GROUP BY")
			}
			q.GroupBy = append(q.GroupBy, t.s)
			if !p.acceptPunct(",") {
				break
			}
		}
	}

	if p.acceptKeyword("order") {
		if err := p.expectKeyword("by"); err != nil {
			return nil, err
		}
		expr, err := p.parseOrderExpr()
		if err != nil {
			return nil, err
		}
		q.OrderBy = expr
		if p.acceptKeyword("desc") {
			q.Desc = true
		} else {
			p.acceptKeyword("asc")
		}
	}

	if p.acceptKeyword("limit") {
		t, ok := p.next()
		if !ok || t.kind != tokNumber {
			return nil, fmt.Errorf("expected a number after LIMIT")
		}
		n, err := strconv.Atoi(t.s)
		if err != nil || n < 0 {
			return nil, fmt.Errorf("LIMIT must be a non-negative integer, got %q", t.s)
		}
		q.Limit = n
	}

	if t, ok := p.peek(); ok {
		return nil, fmt.Errorf("unexpected trailing input %q at position %d", t.s, t.pos)
	}
	if err := q.validate(); err != nil {
		return nil, err
	}
	return q, nil
}

func (p *parser) parseSelects() ([]Select, error) {
	var out []Select
	for {
		s, err := p.parseSelect()
		if err != nil {
			return nil, err
		}
		out = append(out, s)
		if !p.acceptPunct(",") {
			break
		}
	}
	return out, nil
}

func (p *parser) parseSelect() (Select, error) {
	if p.acceptPunct("*") {
		return Select{Column: "*"}, nil
	}
	t, ok := p.next()
	if !ok || t.kind != tokIdent {
		return Select{}, fmt.Errorf("expected a column or aggregate in SELECT")
	}
	// An identifier followed by "(" is an aggregate call.
	if p.acceptPunct("(") {
		agg, err := parseAgg(t.s)
		if err != nil {
			return Select{}, err
		}
		var col string
		if p.acceptPunct("*") {
			col = "*"
		} else {
			ct, ok := p.next()
			if !ok || ct.kind != tokIdent {
				return Select{}, fmt.Errorf("expected a column inside %s(...)", t.s)
			}
			col = ct.s
		}
		if !p.acceptPunct(")") {
			return Select{}, fmt.Errorf("missing closing parenthesis after %s(%s", t.s, col)
		}
		if col == "*" && agg != AggCount {
			return Select{}, fmt.Errorf("%s(*) is not meaningful; only count(*) is", agg)
		}
		sel := Select{Agg: agg, Column: col}
		sel.Alias = p.parseAlias()
		return sel, nil
	}
	sel := Select{Column: t.s}
	sel.Alias = p.parseAlias()
	return sel, nil
}

func (p *parser) parseAlias() string {
	save := p.i
	if p.acceptKeyword("as") {
		if t, ok := p.next(); ok && t.kind == tokIdent {
			return t.s
		}
		p.i = save
	}
	return ""
}

func parseAgg(name string) (AggFunc, error) {
	switch strings.ToLower(name) {
	case "count":
		return AggCount, nil
	case "sum":
		return AggSum, nil
	case "avg":
		return AggAvg, nil
	case "min":
		return AggMin, nil
	case "max":
		return AggMax, nil
	case "p50", "median":
		return AggP50, nil
	case "p90":
		return AggP90, nil
	case "p95":
		return AggP95, nil
	case "p99":
		return AggP99, nil
	}
	return AggNone, fmt.Errorf("unknown aggregate %q; supported: count sum avg min max p50 p90 p95 p99", name)
}

func (p *parser) parseOrderExpr() (string, error) {
	t, ok := p.next()
	if !ok || t.kind != tokIdent {
		return "", fmt.Errorf("expected a column or aggregate in ORDER BY")
	}
	if p.acceptPunct("(") {
		var col string
		if p.acceptPunct("*") {
			col = "*"
		} else {
			ct, ok := p.next()
			if !ok || ct.kind != tokIdent {
				return "", fmt.Errorf("expected a column inside ORDER BY %s(...)", t.s)
			}
			col = ct.s
		}
		if !p.acceptPunct(")") {
			return "", fmt.Errorf("missing closing parenthesis in ORDER BY")
		}
		return strings.ToLower(t.s) + "(" + col + ")", nil
	}
	return t.s, nil
}

func (p *parser) parseWhere(q *Query) error {
	for {
		c, err := p.parseCondition()
		if err != nil {
			return err
		}
		// `time` is not a stored column; it is the query-level range that
		// selects which shards to open. Hoisting it here is what makes the
		// range mandatory and unmissable.
		if c.Column == "time" {
			if !c.IsNum {
				return fmt.Errorf("the `time` bound must be a numeric timestamp in unix micros, got %q", c.Value)
			}
			v := int64(c.Num)
			switch c.Op {
			case OpGe, OpGt:
				q.MinTS = v
			case OpLe, OpLt:
				q.MaxTS = v
			case OpEq:
				q.MinTS, q.MaxTS = v, v
			default:
				return fmt.Errorf("`time` supports >=, >, <=, <, = — not %s", c.Op)
			}
		} else {
			q.Conds = append(q.Conds, c)
		}

		if !p.acceptKeyword("and") {
			break
		}
	}
	if p.acceptKeyword("or") {
		return fmt.Errorf("OR is not supported; conditions are ANDed")
	}
	return nil
}

func (p *parser) parseCondition() (Condition, error) {
	t, ok := p.next()
	if !ok || t.kind != tokIdent {
		return Condition{}, fmt.Errorf("expected a column name in WHERE")
	}
	c := Condition{Column: t.s}

	op, ok := p.next()
	if !ok {
		return Condition{}, fmt.Errorf("expected an operator after %q", c.Column)
	}
	switch {
	case op.kind == tokPunct:
		switch op.s {
		case "=", "!=", "<", "<=", ">", ">=":
			c.Op = Op(op.s)
		default:
			return Condition{}, fmt.Errorf("unsupported operator %q", op.s)
		}
	case op.kind == tokIdent && strings.EqualFold(op.s, "like"):
		c.Op = OpLike
	case op.kind == tokIdent && strings.EqualFold(op.s, "contains"):
		c.Op = OpContains
	default:
		return Condition{}, fmt.Errorf("unsupported operator %q after %q", op.s, c.Column)
	}

	v, ok := p.next()
	if !ok {
		return Condition{}, fmt.Errorf("expected a value after %s %s", c.Column, c.Op)
	}
	switch v.kind {
	case tokString:
		c.Value = v.s
	case tokNumber:
		c.Value = v.s
		if f, err := strconv.ParseFloat(v.s, 64); err == nil {
			c.Num, c.IsNum = f, true
		}
	case tokIdent:
		// Bare words are literals: status = error, kind = tool_call.
		c.Value = v.s
	default:
		return Condition{}, fmt.Errorf("expected a value after %s %s, got %q", c.Column, c.Op, v.s)
	}
	return c, nil
}

func (q *Query) validate() error {
	if q.HasAggregates() {
		// Every non-aggregate projection must be grouped, or the result is a
		// row picked arbitrarily from the group — the classic silently-wrong
		// SQL answer.
		grouped := make(map[string]bool, len(q.GroupBy))
		for _, g := range q.GroupBy {
			grouped[g] = true
		}
		for _, s := range q.Selects {
			if s.Agg == AggNone && !grouped[s.Column] {
				return fmt.Errorf("column %q is selected without an aggregate and is not in GROUP BY", s.Column)
			}
		}
	} else if len(q.GroupBy) > 0 {
		return fmt.Errorf("GROUP BY requires at least one aggregate in SELECT")
	}
	return nil
}