query-interpreter/q/parse.go

package q

import (
	"fmt"
	"strings"

	"github.com/DataDog/go-sqllexer"
)

type Token = sqllexer.Token

type parser struct {
	sql              string
	tokens           []Token
	index            int
	query            Query
	lookBehindBuffer []Token
}

func (p *parser) lookAhead(count int) Token {
	return p.tokens[p.index+count]
}

func (p *parser) lookBehind(count int) Token {
	return p.tokens[p.index-count]
}

// Returns pointer of the first found Token and its index in parser.Tokens
// Returns -1 as the index if not found
func (p *parser) findToken(condition func(Token) bool) (*Token, int) {
	for i, element := range p.tokens {
		if condition(element) {
			return &element, i
		}
	}
	return nil, -1
}

func filter[T any](items []T, fn func(item T) bool) []T {
	filteredItems := []T{}
	for _, value := range items {
		if fn(value) {
			filteredItems = append(filteredItems, value)
		}
	}
	return filteredItems
}

func unshiftBuffer(buf *[10]sqllexer.Token, value sqllexer.Token) {
	for i := 9; i >= 1; i-- {
		buf[i] = buf[i-1]
	}

	buf[0] = value
}

func Parse(sql string) (Query, error) {
	lexer := sqllexer.New(sql)

	var tokens []Token
	for {
		t := lexer.Scan()
		tokens = append(tokens, *t)

		if IsTokenEndOfStatement(t) {
			break
		}
	}

	p := parser{
		sql:    sql,
		tokens: tokens,
	}

	queryType, err := findQueryType(&p)
	if err != nil {
		return p.query, err
	}

	switch strings.ToUpper(queryType) {
	case "SELECT":
		p.query = &Select{
			Type: SELECT,
		}
		parseSelectStatement(&p)
		fmt.Printf("IsWildcard %v", p.query.(*Select).IsWildcard)
	default:
		return p.query, fmt.Errorf("No process defined for determined queryType: %s", queryType)
	}

	return p.query, nil
}

func findQueryType(p *parser) (string, error) {
	firstCommand, _ := p.findToken(func(t Token) bool {
		return t.Type == sqllexer.COMMAND
	})

	if firstCommand == nil {
		return "", fmt.Errorf("Could not find query type")
	}

	return firstCommand.Value, nil
}

func parseSelectStatement(p *parser) {
	selectQuery := p.query.(*Select)

	foundWildcard, _ := p.findToken(func(t Token) bool {
		return t.Type == sqllexer.WILDCARD
	})

	selectQuery.IsWildcard = foundWildcard != nil

	if !selectQuery.IsWildcard {
		parseSelectColumns(p)
	}

}

func parseSelectColumns(p *parser) error {
	selectQuery := p.query.(*Select)

	_, selectCommandIndex := p.findToken(func(t Token) bool {
		return t.Type == sqllexer.COMMAND && strings.ToUpper(t.Value) == "SELECT"
	})

	_, fromKeywordIndex := p.findToken(func(t Token) bool {
		return t.Type == sqllexer.KEYWORD && strings.ToUpper(t.Value) == "FROM"
	})

	if selectCommandIndex < 0 || fromKeywordIndex < 0 {
		return fmt.Errorf("Could not find range between SELECT and FROM")
	}

	tokensFromSelectToFrom := p.tokens[selectCommandIndex:fromKeywordIndex]

	lookBehindBuffer := [10]Token{}
	var workingColumn Column
	columns := make([]Column, 0)

	for _, token := range tokensFromSelectToFrom {
		if token.Type == sqllexer.FUNCTION {
			unshiftBuffer(&lookBehindBuffer, token)
			workingColumn.AggregateFunction = AggregateFunctionTypeByName(token.Value)
			continue
		} else if token.Type == sqllexer.PUNCTUATION {
			if token.Value == "," {
				columns = append(columns, workingColumn)
				workingColumn = Column{}
			}
			continue
		} else if token.Type == sqllexer.IDENT {
			unshiftBuffer(&lookBehindBuffer, token)

			if lookBehindBuffer[0].Type == sqllexer.ALIAS_INDICATOR {
				workingColumn.Alias = token.Value
			} else {
				workingColumn.Name = token.Value
			}

			continue
		} else if token.Type == sqllexer.ALIAS_INDICATOR {
			unshiftBuffer(&lookBehindBuffer, token)
			continue
		} else if token.Type == sqllexer.SPACE {
			continue
		} else {
			if workingColumn.Name != "" {
				columns = append(columns, workingColumn)
				workingColumn = Column{}
			}
		}
	}

	selectQuery.Columns = columns

	return nil
}