Skip to content

Commit

Permalink
Base model.Expr visitor (#531)
Browse files Browse the repository at this point in the history 
We introduce a base visitor that implements the `model.ExprVisitor`. We
usually implement only one or two methods for the visitor. Other methods
are the same.

@pdelewski please take a look at this
  • Loading branch information
@nablaone
nablaone authored Jul 16, 2024
1 parent b5a5e9d commit a6de33f
Show file tree
Hide file tree
Showing 9 changed files with 460 additions and 909 deletions.
200 changes: 200 additions & 0 deletions quesma/model/base_visitor.go
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,200 @@
// Copyright Quesma, licensed under the Elastic License 2.0.
// SPDX-License-Identifier: Elastic-2.0
package model

type BaseExprVisitor struct {
OverrideVisitFunction func(b *BaseExprVisitor, e FunctionExpr) interface{}
OverrideVisitMultiFunction func(b *BaseExprVisitor, e MultiFunctionExpr) interface{}
OverrideVisitLiteral func(b *BaseExprVisitor, l LiteralExpr) interface{}
OverrideVisitString func(b *BaseExprVisitor, e StringExpr) interface{}
OverrideVisitInfix func(b *BaseExprVisitor, e InfixExpr) interface{}
OverrideVisitColumnRef func(b *BaseExprVisitor, e ColumnRef) interface{}
OverrideVisitPrefixExpr func(b *BaseExprVisitor, e PrefixExpr) interface{}
OverrideVisitNestedProperty func(b *BaseExprVisitor, e NestedProperty) interface{}
OverrideVisitArrayAccess func(b *BaseExprVisitor, e ArrayAccess) interface{}
OverrideVisitOrderByExpr func(b *BaseExprVisitor, e OrderByExpr) interface{}
OverrideVisitDistinctExpr func(b *BaseExprVisitor, e DistinctExpr) interface{}
OverrideVisitTableRef func(b *BaseExprVisitor, e TableRef) interface{}
OverrideVisitAliasedExpr func(b *BaseExprVisitor, e AliasedExpr) interface{}
OverrideVisitSelectCommand func(b *BaseExprVisitor, e SelectCommand) interface{}
OverrideVisitWindowFunction func(b *BaseExprVisitor, f WindowFunction) interface{}
OverrideVisitParenExpr func(b *BaseExprVisitor, e ParenExpr) interface{}
OverrideVisitLambdaExpr func(b *BaseExprVisitor, e LambdaExpr) interface{}
}

func NewBaseVisitor() *BaseExprVisitor {
return &BaseExprVisitor{}
}

func (v *BaseExprVisitor) VisitChildren(args []Expr) []Expr {
var newArgs []Expr
for _, arg := range args {
if arg != nil {
newArgs = append(newArgs, arg.Accept(v).(Expr))
}
}
return newArgs
}

func (v *BaseExprVisitor) VisitLiteral(e LiteralExpr) interface{} {
if v.OverrideVisitLiteral != nil {
return v.OverrideVisitLiteral(v, e)
}

return NewLiteral(e.Value)
}
func (v *BaseExprVisitor) VisitInfix(e InfixExpr) interface{} {
if v.OverrideVisitInfix != nil {
return v.OverrideVisitInfix(v, e)
}
lhs := e.Left.Accept(v)
rhs := e.Right.Accept(v)
return NewInfixExpr(lhs.(Expr), e.Op, rhs.(Expr))
}

func (v *BaseExprVisitor) VisitPrefixExpr(e PrefixExpr) interface{} {
if v.OverrideVisitPrefixExpr != nil {
return v.OverrideVisitPrefixExpr(v, e)
}
return NewPrefixExpr(e.Op, v.VisitChildren(e.Args))
}

func (v *BaseExprVisitor) VisitFunction(e FunctionExpr) interface{} {
if v.OverrideVisitFunction != nil {
return v.OverrideVisitFunction(v, e)
}
return NewFunction(e.Name, v.VisitChildren(e.Args)...)
}

func (v *BaseExprVisitor) VisitColumnRef(e ColumnRef) interface{} {
if v.OverrideVisitColumnRef != nil {
return v.OverrideVisitColumnRef(v, e)
}
return NewColumnRef(e.ColumnName)
}

func (v *BaseExprVisitor) VisitNestedProperty(e NestedProperty) interface{} {
if v.OverrideVisitNestedProperty != nil {
return v.OverrideVisitNestedProperty(v, e)
}
ColumnRef := e.ColumnRef.Accept(v).(ColumnRef)
Property := e.PropertyName.Accept(v).(LiteralExpr)
return NewNestedProperty(ColumnRef, Property)
}

func (v *BaseExprVisitor) VisitArrayAccess(e ArrayAccess) interface{} {
if v.OverrideVisitArrayAccess != nil {
return v.OverrideVisitArrayAccess(v, e)
}
columnRef := e.ColumnRef.Accept(v).(ColumnRef)
index := e.Index.Accept(v).(Expr)
return NewArrayAccess(columnRef, index)
}

func (v *BaseExprVisitor) VisitMultiFunction(e MultiFunctionExpr) interface{} {
if v.OverrideVisitMultiFunction != nil {
return v.OverrideVisitMultiFunction(v, e)
}
return MultiFunctionExpr{Name: e.Name, Args: v.VisitChildren(e.Args)}
}

func (v *BaseExprVisitor) VisitString(e StringExpr) interface{} {
if v.OverrideVisitString != nil {
return v.OverrideVisitString(v, e)
}
return e
}

func (v *BaseExprVisitor) VisitTableRef(e TableRef) interface{} {
if v.OverrideVisitTableRef != nil {
return v.OverrideVisitTableRef(v, e)
}
return e
}

func (v *BaseExprVisitor) VisitOrderByExpr(e OrderByExpr) interface{} {
if v.OverrideVisitOrderByExpr != nil {
return v.OverrideVisitOrderByExpr(v, e)
}
return OrderByExpr{Exprs: v.VisitChildren(e.Exprs), Direction: e.Direction}
}

func (v *BaseExprVisitor) VisitDistinctExpr(e DistinctExpr) interface{} {
if v.OverrideVisitDistinctExpr != nil {
return v.OverrideVisitDistinctExpr(v, e)
}
return DistinctExpr{Expr: e.Expr.Accept(v).(Expr)}
}

func (v *BaseExprVisitor) VisitAliasedExpr(e AliasedExpr) interface{} {
if v.OverrideVisitAliasedExpr != nil {
return v.OverrideVisitAliasedExpr(v, e)
}
return NewAliasedExpr(e.Expr.Accept(v).(Expr), e.Alias)
}

func (v *BaseExprVisitor) VisitWindowFunction(f WindowFunction) interface{} {
if v.OverrideVisitWindowFunction != nil {
return v.OverrideVisitWindowFunction(v, f)
}
return WindowFunction{
Name: f.Name,
Args: v.VisitChildren(f.Args),
PartitionBy: v.VisitChildren(f.PartitionBy),
OrderBy: f.OrderBy.Accept(v).(OrderByExpr),
}
}

func (v *BaseExprVisitor) VisitSelectCommand(query SelectCommand) interface{} {
if v.OverrideVisitSelectCommand != nil {
return v.OverrideVisitSelectCommand(v, query)
}
var columns, groupBy []Expr
var orderBy []OrderByExpr
from := query.FromClause
where := query.WhereClause

for _, expr := range query.Columns {
columns = append(columns, expr.Accept(v).(Expr))
}
for _, expr := range query.GroupBy {
groupBy = append(groupBy, expr.Accept(v).(Expr))
}
for _, expr := range query.OrderBy {
orderBy = append(orderBy, expr.Accept(v).(OrderByExpr))
}
if query.FromClause != nil {
from = query.FromClause.Accept(v).(Expr)
}
if query.WhereClause != nil {
where = query.WhereClause.Accept(v).(Expr)
}

var ctes []*SelectCommand
if query.CTEs != nil {
ctes = make([]*SelectCommand, 0)
for _, cte := range query.CTEs {
ctes = append(ctes, cte.Accept(v).(*SelectCommand))
}
}

return NewSelectCommand(columns, groupBy, orderBy, from, where, query.LimitBy, query.Limit, query.SampleLimit, query.IsDistinct, ctes)
}

func (v *BaseExprVisitor) VisitParenExpr(p ParenExpr) interface{} {
if v.OverrideVisitParenExpr != nil {
return v.OverrideVisitParenExpr(v, p)
}
var exprs []Expr
for _, expr := range p.Exprs {
exprs = append(exprs, expr.Accept(v).(Expr))
}
return NewParenExpr(exprs...)
}

func (v *BaseExprVisitor) VisitLambdaExpr(e LambdaExpr) interface{} {
if v.OverrideVisitLambdaExpr != nil {
return v.OverrideVisitLambdaExpr(v, e)
}
return NewLambdaExpr(e.Args, e.Body.Accept(v).(Expr))
}
180 changes: 31 additions & 149 deletions quesma/model/highlighter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,9 +50,37 @@ func (h *Highlighter) ShouldHighlight(columnName string) bool {

// SetTokensToHighlight takes a Select query and extracts tokens that should be highlighted.
func (h *Highlighter) SetTokensToHighlight(selectCmd SelectCommand) {
highlighterVisitor := NewHighlighter()
selectCmd.Accept(highlighterVisitor)
h.Tokens = highlighterVisitor.Tokens

h.Tokens = make(map[string]Tokens)

visitor := NewBaseVisitor()

visitor.OverrideVisitInfix = func(b *BaseExprVisitor, e InfixExpr) interface{} {
switch e.Op {
case "iLIKE", "LIKE", "IN", "=":
lhs, isColumnRef := e.Left.(ColumnRef)
rhs, isLiteral := e.Right.(LiteralExpr)
if isLiteral && isColumnRef { // we only highlight in this case
switch literalAsString := rhs.Value.(type) {
case string:
literalAsString = strings.TrimPrefix(literalAsString, "'%")
literalAsString = strings.TrimPrefix(literalAsString, "%")
literalAsString = strings.TrimSuffix(literalAsString, "'")
literalAsString = strings.TrimSuffix(literalAsString, "%")
if h.Tokens[lhs.ColumnName] == nil {
h.Tokens[lhs.ColumnName] = make(Tokens)
}
h.Tokens[lhs.ColumnName][strings.ToLower(literalAsString)] = struct{}{}
default:
logger.Info().Msgf("Value is of an unexpected type: %T\n", literalAsString)
}
}
}
return NewInfixExpr(e.Left.Accept(b).(Expr), e.Op, e.Right.Accept(b).(Expr))
}

selectCmd.Accept(visitor)

}

// HighlightValue takes a value and returns the part of it that should be highlighted, wrapped in tags.
Expand Down Expand Up @@ -124,149 +152,3 @@ func (h *Highlighter) HighlightValue(columnName, value string) []string {

return highlights
}

// highlighter is a visitor that traverses the AST and collects tokens that should be highlighted.
type highlighter struct {
// TokensToHighlight represents a set of tokens that should be highlighted in the query.
Tokens map[string]Tokens
}

func NewHighlighter() *highlighter {
return &highlighter{
Tokens: make(map[string]Tokens),
}
}

func (v *highlighter) VisitColumnRef(e ColumnRef) interface{} {
return e
}

func (v *highlighter) VisitPrefixExpr(e PrefixExpr) interface{} {
var exprs []Expr
for _, expr := range e.Args {
exprs = append(exprs, expr.Accept(v).(Expr))
}
return NewPrefixExpr(e.Op, exprs)
}

func (v *highlighter) VisitNestedProperty(e NestedProperty) interface{} {
return NewNestedProperty(e.ColumnRef.Accept(v).(ColumnRef), e.PropertyName)
}

func (v *highlighter) VisitArrayAccess(e ArrayAccess) interface{} {
return NewArrayAccess(e.ColumnRef.Accept(v).(ColumnRef), e.Index.Accept(v).(Expr))
}

func (v *highlighter) VisitFunction(e FunctionExpr) interface{} {
var exprs []Expr
for _, expr := range e.Args {
exprs = append(exprs, expr.Accept(v).(Expr))
}
return NewFunction(e.Name, exprs...)
}

func (v *highlighter) VisitLiteral(l LiteralExpr) interface{} {
return l
}

func (v *highlighter) VisitString(e StringExpr) interface{} {
return e
}

func (v *highlighter) VisitMultiFunction(f MultiFunctionExpr) interface{} {
var exprs []Expr
for _, expr := range f.Args {
exprs = append(exprs, expr.Accept(v).(Expr))
}
return MultiFunctionExpr{Name: f.Name, Args: exprs}
}

func (v *highlighter) VisitInfix(e InfixExpr) interface{} {
switch e.Op {
case "iLIKE", "LIKE", "IN", "=":
lhs, isColumnRef := e.Left.(ColumnRef)
rhs, isLiteral := e.Right.(LiteralExpr)
if isLiteral && isColumnRef { // we only highlight in this case
switch literalAsString := rhs.Value.(type) {
case string:
literalAsString = strings.TrimPrefix(literalAsString, "'%")
literalAsString = strings.TrimPrefix(literalAsString, "%")
literalAsString = strings.TrimSuffix(literalAsString, "'")
literalAsString = strings.TrimSuffix(literalAsString, "%")
if v.Tokens[lhs.ColumnName] == nil {
v.Tokens[lhs.ColumnName] = make(Tokens)
}
v.Tokens[lhs.ColumnName][strings.ToLower(literalAsString)] = struct{}{}
default:
logger.Info().Msgf("Value is of an unexpected type: %T\n", literalAsString)
}
}
}
return NewInfixExpr(e.Left.Accept(v).(Expr), e.Op, e.Right.Accept(v).(Expr))
}

func (v *highlighter) VisitOrderByExpr(e OrderByExpr) interface{} {
var exprs []Expr
for _, expr := range e.Exprs {
exprs = append(exprs, expr.Accept(v).(Expr))
}
return NewOrderByExpr(exprs, e.Direction)
}

func (v *highlighter) VisitDistinctExpr(e DistinctExpr) interface{} {
return NewDistinctExpr(e.Expr.Accept(v).(Expr))
}

func (v *highlighter) VisitTableRef(e TableRef) interface{} {
return e
}

func (v *highlighter) VisitAliasedExpr(e AliasedExpr) interface{} {
return NewAliasedExpr(e.Expr.Accept(v).(Expr), e.Alias)
}

func (v *highlighter) VisitSelectCommand(c SelectCommand) interface{} {
var columns, groupBy []Expr
var orderBy []OrderByExpr
from := c.FromClause
where := c.WhereClause
for _, expr := range c.Columns {
columns = append(columns, expr.Accept(v).(Expr))
}
for _, expr := range c.GroupBy {
groupBy = append(groupBy, expr.Accept(v).(Expr))
}
for _, expr := range c.OrderBy {
orderBy = append(orderBy, expr.Accept(v).(OrderByExpr))
}
if c.FromClause != nil {
from = c.FromClause.Accept(v).(Expr)
}
if c.WhereClause != nil {
where = c.WhereClause.Accept(v).(Expr)
}
return *NewSelectCommand(columns, groupBy, orderBy, from, where, []Expr{}, c.Limit, c.SampleLimit, c.IsDistinct, c.CTEs)
}

func (v *highlighter) VisitWindowFunction(f WindowFunction) interface{} {
var args, partitionBy []Expr
for _, expr := range f.Args {
args = append(args, expr.Accept(v).(Expr))
}
for _, expr := range f.PartitionBy {
partitionBy = append(partitionBy, expr.Accept(v).(Expr))
}
return NewWindowFunction(f.Name, args, partitionBy, f.OrderBy.Accept(v).(OrderByExpr))
}

func (v *highlighter) VisitParenExpr(p ParenExpr) interface{} {
var exprs []Expr
for _, expr := range p.Exprs {
exprs = append(exprs, expr.Accept(v).(Expr))
}
return NewParenExpr(exprs...)
}

func (v *highlighter) VisitLambdaExpr(l LambdaExpr) interface{} {
return NewLambdaExpr(l.Args, l.Body.Accept(v).(Expr))
}
Loading

0 comments on commit a6de33f

Please sign in to comment.