forked from vitessio/vitess
-
Notifications
You must be signed in to change notification settings - Fork 0
/
normalizer.go
131 lines (124 loc) · 3.61 KB
/
normalizer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
// Copyright 2016, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package sqlparser
import (
"fmt"
"github.com/youtube/vitess/go/sqltypes"
querypb "github.com/youtube/vitess/go/vt/proto/query"
)
// Normalize changes the statement to use bind values, and
// updates the bind vars to those values. The supplied prefix
// is used to generate the bind var names. The function ensures
// that there are no collisions with existing bind vars.
func Normalize(stmt Statement, bindVars map[string]interface{}, prefix string) {
reserved := GetBindvars(stmt)
// vals allows us to reuse bindvars for
// identical values.
counter := 1
vals := make(map[string]string)
_ = Walk(func(node SQLNode) (kontinue bool, err error) {
switch node := node.(type) {
case *SQLVal:
// Make the bindvar
bval := sqlToBindvar(node)
if bval == nil {
// If unsuccessful continue.
return true, nil
}
// Check if there's a bindvar for that value already.
var key string
if bval.Type == sqltypes.VarBinary {
// Prefixing strings with "'" ensures that a string
// and number that have the same representation don't
// collide.
key = "'" + string(node.Val)
} else {
key = string(node.Val)
}
bvname, ok := vals[key]
if !ok {
// If there's no such bindvar, make a new one.
bvname, counter = newName(prefix, counter, reserved)
vals[key] = bvname
bindVars[bvname] = bval
}
// Modify the AST node to a bindvar.
node.Type = ValArg
node.Val = append([]byte(":"), bvname...)
case *ComparisonExpr:
switch node.Operator {
case InStr, NotInStr:
default:
return true, nil
}
// It's either IN or NOT IN.
tupleVals, ok := node.Right.(ValTuple)
if !ok {
return true, nil
}
// The RHS is a tuple of values.
// Make a list bindvar.
bvals := &querypb.BindVariable{
Type: sqltypes.Tuple,
}
for _, val := range tupleVals {
bval := sqlToBindvar(val)
if bval == nil {
return true, nil
}
bvals.Values = append(bvals.Values, &querypb.Value{
Type: bval.Type,
Value: bval.Value,
})
}
var bvname string
bvname, counter = newName(prefix, counter, reserved)
bindVars[bvname] = bvals
// Modify RHS to be a list bindvar.
node.Right = ListArg(append([]byte("::"), bvname...))
}
return true, nil
}, stmt)
}
func sqlToBindvar(node SQLNode) *querypb.BindVariable {
if node, ok := node.(*SQLVal); ok {
switch node.Type {
case StrVal:
return &querypb.BindVariable{Type: sqltypes.VarBinary, Value: node.Val}
case IntVal:
return &querypb.BindVariable{Type: sqltypes.Int64, Value: node.Val}
case FloatVal:
return &querypb.BindVariable{Type: sqltypes.Float64, Value: node.Val}
}
}
return nil
}
func newName(prefix string, counter int, reserved map[string]struct{}) (string, int) {
for {
newName := fmt.Sprintf("%s%d", prefix, counter)
if _, ok := reserved[newName]; !ok {
reserved[newName] = struct{}{}
return newName, counter + 1
}
counter++
}
}
// GetBindvars returns a map of the bind vars referenced in the statement.
// TODO(sougou); This function gets called again from vtgate/planbuilder.
// Ideally, this should be done only once.
func GetBindvars(stmt Statement) map[string]struct{} {
bindvars := make(map[string]struct{})
_ = Walk(func(node SQLNode) (kontinue bool, err error) {
switch node := node.(type) {
case *SQLVal:
if node.Type == ValArg {
bindvars[string(node.Val[1:])] = struct{}{}
}
case ListArg:
bindvars[string(node[2:])] = struct{}{}
}
return true, nil
}, stmt)
return bindvars
}