84 lines
2.5 KiB
Go

package main
import (
"fmt"
"github.com/DataDog/go-sqllexer"
)
type TokenGroupType int
const (
TOKEN_OTHER TokenGroupType = iota
TOKEN_ACT // SELECT, CREATE, etc
TOKEN_FILTER // WHERE, AND, etc
TOKEN_RELATE // joins
TOKEN_DEFINE // column data types, etc
TOKEN_CONTROL // sorting, etc
)
var tokenGroupMap = map[sqllexer.TokenType]TokenGroupType{
sqllexer.ERROR: TOKEN_OTHER,
sqllexer.EOF: TOKEN_OTHER,
sqllexer.SPACE: TOKEN_OTHER,
sqllexer.STRING: TOKEN_DEFINE,
sqllexer.INCOMPLETE_STRING: TOKEN_DEFINE,
sqllexer.NUMBER: TOKEN_DEFINE,
sqllexer.IDENT: TOKEN_ACT, // for table/column names
sqllexer.QUOTED_IDENT: TOKEN_RELATE, // for joins
sqllexer.OPERATOR: TOKEN_FILTER,
sqllexer.WILDCARD: TOKEN_ACT,
sqllexer.COMMENT: TOKEN_OTHER,
sqllexer.MULTILINE_COMMENT: TOKEN_OTHER,
sqllexer.PUNCTUATION: TOKEN_OTHER,
sqllexer.DOLLAR_QUOTED_FUNCTION: TOKEN_DEFINE,
sqllexer.DOLLAR_QUOTED_STRING: TOKEN_DEFINE,
sqllexer.POSITIONAL_PARAMETER: TOKEN_FILTER,
sqllexer.BIND_PARAMETER: TOKEN_FILTER,
sqllexer.FUNCTION: TOKEN_DEFINE,
sqllexer.SYSTEM_VARIABLE: TOKEN_DEFINE,
sqllexer.UNKNOWN: TOKEN_OTHER,
sqllexer.COMMAND: TOKEN_ACT,
sqllexer.KEYWORD: TOKEN_ACT,
sqllexer.JSON_OP: TOKEN_DEFINE,
sqllexer.BOOLEAN: TOKEN_DEFINE,
sqllexer.NULL: TOKEN_DEFINE,
sqllexer.PROC_INDICATOR: TOKEN_OTHER,
sqllexer.CTE_INDICATOR: TOKEN_OTHER,
sqllexer.ALIAS_INDICATOR: TOKEN_OTHER,
}
func GetTokenGroupType(tokenType sqllexer.TokenType) (TokenGroupType, bool, bool) {
group, exists := tokenGroupMap[tokenType]
isStart := false
isEnd := false
if !exists {
group = TOKEN_OTHER
}
if group == TOKEN_ACT {
isStart = true
}
if tokenType == sqllexer.EOF ||
(group == TOKEN_OTHER && tokenType == sqllexer.PUNCTUATION) {
isEnd = true
}
return group, isStart, isEnd
}
func main() {
query := "SELECT * FROM users \n WHERE id = something AND SELECT;"
lexer := sqllexer.New(query)
for {
token := lexer.Scan()
tokenGroup, isStart, isEnd := GetTokenGroupType(token.Type)
fmt.Println(token.Value, token.Type, tokenGroup, isStart, isEnd)
if token.Type == sqllexer.EOF {
break
}
}
}