forked from ruiaylin/pgparser
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pretty_scanner.go
169 lines (158 loc) · 5.33 KB
/
pretty_scanner.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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
// Copyright 2019 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.
package keysutils
import (
"fmt"
"strings"
"github.com/ruiaylin/pgparser/pkg/keys"
"github.com/ruiaylin/pgparser/pkg/roachpb"
"github.com/ruiaylin/pgparser/ast"
"github.com/ruiaylin/pgparser/sqlbase"
"github.com/ruiaylin/pgparser/utils/encoding"
"github.com/ruiaylin/pgparser/utils/keysutil"
)
// MakePrettyScannerForNamedTables create a PrettyScanner that, beside what the
// PrettyScanner is generally able to decode, can also decode keys of the form
// "/<table name>/<index name>/1/2/3/..." using supplied maps from names to ids.
//
// TODO(nvanbenschoten): support tenant SQL keys.
func MakePrettyScannerForNamedTables(
tableNameToID map[string]int, idxNameToID map[string]int,
) keysutil.PrettyScanner {
return keysutil.MakePrettyScanner(func(input string) (string, roachpb.Key) {
remainder, k := parseTableKeysAsAscendingInts(input, tableNameToID, idxNameToID)
return remainder, k
})
}
// parseTableKeysAsAscendingInts takes a pretty-printed key segment like
// "/<table name>/<index name>/1/2/3/...", a mapping of table names to ids and a
// mapping of index names to ids and turns the part before "/... " into the key
// corresponding to the respective index entry by encoding the ints as ascending
// varints (so assuming that they correspond to columns of various integer
// types).
//
// idxNameToID has to contain entries of the form "<table name>.<index name>".
// The index name "pk" is implicitly mapped to index 1; it doesn't need to
// appear in idxNameToID.
//
// Notice that the input is not expected to have the "/Table" prefix (which is
// generated by pretty-printing keys). That prefix is assumed to have been
// consumed before this function is invoked (by the PrettyScanner).
//
// The "/..." part is returned as the remainder (the first return value).
func parseTableKeysAsAscendingInts(
input string, tableNameToID map[string]int, idxNameToID map[string]int,
) (string, roachpb.Key) {
// Consume the table name.
input = mustShiftSlash(input)
slashPos := strings.Index(input, "/")
if slashPos < 0 {
slashPos = len(input)
}
remainder := input[slashPos:] // `/something/else` -> `/else`
tableName := input[:slashPos]
tableID, ok := tableNameToID[tableName]
if !ok {
panic(fmt.Sprintf("unknown table: %s", tableName))
}
output := keys.TODOSQLCodec.TablePrefix(uint32(tableID))
if remainder == "" {
return "", output
}
input = remainder
// Consume the index name.
input = mustShiftSlash(input)
slashPos = strings.Index(input, "/")
if slashPos < 0 {
// Accept a string ending in "/<index name>.
slashPos = len(input)
}
remainder = input[slashPos:] // `/something/else` -> `/else`
idxName := input[:slashPos]
var idxID int
// The primary key index always has ID 1.
if idxName == "pk" {
idxID = 1
} else {
idxID, ok = idxNameToID[fmt.Sprintf("%s.%s", tableName, idxName)]
if !ok {
panic(fmt.Sprintf("unknown index: %s", idxName))
}
}
output = encoding.EncodeUvarintAscending(output, uint64(idxID))
if remainder == "" {
return "", output
}
input = remainder
remainder, moreOutput := parseAscendingIntIndexKeys(input)
output = append(output, moreOutput...)
return remainder, output
}
func mustShiftSlash(in string) string {
slash, out := mustShift(in)
if slash != "/" {
panic("expected /: " + in)
}
return out
}
func mustShift(in string) (first, remainder string) {
if len(in) == 0 {
panic("premature end of string")
}
return in[:1], in[1:]
}
// parseAscendingIntIndexKeys takes a pretty-printed key segment like
// "/1/2/3/foo" and parses all the ints from the beginning turning them into a
// key segment by encoding them as ascending varints. The part after the last
// int is returned as the remainder (the first return value).
func parseAscendingIntIndexKeys(input string) (string, roachpb.Key) {
var key roachpb.Key
for {
remainder, k := parseAscendingIntIndexKey(input)
if k == nil {
// We've failed to parse anything.
return remainder, key
}
key = append(key, k...)
input = remainder
if remainder == "" {
// We've consumed the whole string.
return "", key
}
}
}
// ParseAscendingIntIndexKey parses one int out of a string looking like
// "/1[/...]" and encodes it as an ascending varint. The rest of the input string
// is returned as the remainder (the first return value).
//
// If the beginning of the input can't be parsed, the returned key is nil and
// the whole input is returned as the remainder.
func parseAscendingIntIndexKey(input string) (string, roachpb.Key) {
origInput := input
input = mustShiftSlash(input)
slashPos := strings.Index(input, "/")
if slashPos < 0 {
// Deal with the case where there's no remainder: the entire string is the
// index ID.
slashPos = len(input)
}
indexValStr := input[:slashPos]
datum, err := ast.ParseDInt(indexValStr)
if err != nil {
// Can't decode the key.
return origInput, nil
}
remainder := input[slashPos:] // `/something/else` -> `/else`
key, err := sqlbase.EncodeTableKey(nil, datum, encoding.Ascending)
if err != nil {
panic(err)
}
return remainder, key
}