Skip to content

Commit f26168c

Browse files
mcarmonaamcarmonaa
committed
internal/function: add uast_extract udf
Signed-off-by: Manuel Carmona <[email protected]>
1 parent 96b30e9 commit f26168c

File tree

3 files changed

+261
-32
lines changed

3 files changed

+261
-32
lines changed

‎internal/function/registry.go

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,11 @@ import "gopkg.in/src-d/go-mysql-server.v0/sql"
44

55
// Functions for gitbase queries.
66
var Functions = sql.Functions{
7-
"is_tag": sql.Function1(NewIsTag),
8-
"is_remote": sql.Function1(NewIsRemote),
9-
"language": sql.FunctionN(NewLanguage),
10-
"uast": sql.FunctionN(NewUAST),
11-
"uast_mode": sql.Function3(NewUASTMode),
12-
"uast_xpath": sql.Function2(NewUASTXPath),
7+
"is_tag": sql.Function1(NewIsTag),
8+
"is_remote": sql.Function1(NewIsRemote),
9+
"language": sql.FunctionN(NewLanguage),
10+
"uast": sql.FunctionN(NewUAST),
11+
"uast_mode": sql.Function3(NewUASTMode),
12+
"uast_xpath": sql.Function2(NewUASTXPath),
13+
"uast_extract": sql.Function2(NewUASTExtract),
1314
}

‎internal/function/uast.go

Lines changed: 137 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -275,12 +275,7 @@ func (f UASTMode) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err erro
275275
return nil, fmt.Errorf("invalid uast mode %s", m)
276276
}
277277

278-
u, err := getUAST(ctx, bytes, lang, "", mode)
279-
if err != nil {
280-
return nil, err
281-
}
282-
283-
return u, nil
278+
return getUAST(ctx, bytes, lang, "", mode)
284279
}
285280

286281
// UASTXPath performs an XPath query over the given UAST nodes.
@@ -318,37 +313,20 @@ func (f *UASTXPath) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err er
318313
return nil, nil
319314
}
320315

321-
left, err = sql.Array(sql.Blob).Convert(left)
316+
nodes, err := nodesFromBlobArray(left)
322317
if err != nil {
323318
return nil, err
324319
}
325320

326-
arr := left.([]interface{})
327-
var nodes = make([]*uast.Node, len(arr))
328-
for i, n := range arr {
329-
node := uast.NewNode()
330-
if err := node.Unmarshal(n.([]byte)); err != nil {
331-
return nil, err
332-
}
333-
nodes[i] = node
334-
}
335-
336-
right, err := f.Right.Eval(ctx, row)
321+
xpath, err := exprToString(ctx, f.Right, row)
337322
if err != nil {
338323
return nil, err
339324
}
340325

341-
if right == nil {
326+
if xpath == "" {
342327
return nil, nil
343328
}
344329

345-
right, err = sql.Text.Convert(right)
346-
if err != nil {
347-
return nil, err
348-
}
349-
350-
xpath := right.(string)
351-
352330
var result []interface{}
353331
for _, n := range nodes {
354332
ns, err := tools.Filter(n, xpath)
@@ -368,6 +346,25 @@ func (f *UASTXPath) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err er
368346
return result, nil
369347
}
370348

349+
func nodesFromBlobArray(data interface{}) ([]*uast.Node, error) {
350+
data, err := sql.Array(sql.Blob).Convert(data)
351+
if err != nil {
352+
return nil, err
353+
}
354+
355+
arr := data.([]interface{})
356+
var nodes = make([]*uast.Node, len(arr))
357+
for i, n := range arr {
358+
node := uast.NewNode()
359+
if err := node.Unmarshal(n.([]byte)); err != nil {
360+
return nil, err
361+
}
362+
nodes[i] = node
363+
}
364+
365+
return nodes, nil
366+
}
367+
371368
func (f UASTXPath) String() string {
372369
return fmt.Sprintf("uast_xpath(%s, %s)", f.Left, f.Right)
373370
}
@@ -475,3 +472,117 @@ func getUAST(
475472

476473
return result, nil
477474
}
475+
476+
// UASTExtract extracts keys from an UAST.
477+
type UASTExtract struct {
478+
expression.BinaryExpression
479+
}
480+
481+
// NewUASTExtract creates a new UASTExtract UDF.
482+
func NewUASTExtract(uast, key sql.Expression) sql.Expression {
483+
return &UASTExtract{expression.BinaryExpression{Left: uast, Right: key}}
484+
}
485+
486+
// String implements the fmt.Stringer interface.
487+
func (u *UASTExtract) String() string {
488+
return fmt.Sprintf("uast_extract(%s, %s)", u.Left, u.Right)
489+
}
490+
491+
// Type implements the sql.Expression interface.
492+
func (u *UASTExtract) Type() sql.Type {
493+
return sql.Array(sql.Array(sql.Text))
494+
}
495+
496+
// Eval implements the sql.Expression interface.
497+
func (u *UASTExtract) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err error) {
498+
defer func() {
499+
if r := recover(); r != nil {
500+
err = fmt.Errorf("uast: unknown error: %s", r)
501+
}
502+
}()
503+
504+
span, ctx := ctx.Span("gitbase.UASTExtract")
505+
defer span.Finish()
506+
507+
left, err := u.Left.Eval(ctx, row)
508+
if err != nil {
509+
return nil, err
510+
}
511+
512+
if left == nil {
513+
return nil, nil
514+
}
515+
516+
nodes, err := nodesFromBlobArray(left)
517+
if err != nil {
518+
return nil, err
519+
}
520+
521+
key, err := exprToString(ctx, u.Right, row)
522+
if err != nil {
523+
return nil, err
524+
}
525+
526+
if key == "" {
527+
return nil, nil
528+
}
529+
530+
extracted := make([][]string, len(nodes))
531+
for i, n := range nodes {
532+
extracted[i] = extractInfo(n, key)
533+
}
534+
535+
return extracted, nil
536+
}
537+
538+
const (
539+
keyType = "@type"
540+
keyToken = "@token"
541+
keyRoles = "@role"
542+
keyStartPos = "@startpos"
543+
keyEndPos = "@endpos"
544+
)
545+
546+
func extractInfo(n *uast.Node, key string) []string {
547+
548+
info := []string{}
549+
switch key {
550+
case keyType:
551+
info = append(info, n.InternalType)
552+
case keyToken:
553+
info = append(info, n.Token)
554+
case keyRoles:
555+
roles := make([]string, len(n.Roles))
556+
for i, rol := range n.Roles {
557+
roles[i] = rol.String()
558+
}
559+
560+
info = append(info, roles...)
561+
case keyStartPos:
562+
info = append(info, n.StartPosition.String())
563+
case keyEndPos:
564+
info = append(info, n.EndPosition.String())
565+
default:
566+
v, ok := n.Properties[key]
567+
if ok {
568+
info = append(info, v)
569+
}
570+
}
571+
572+
return info
573+
}
574+
575+
// TransformUp implements the sql.Expression interface.
576+
func (u *UASTExtract) TransformUp(f sql.TransformExprFunc) (sql.Expression, error) {
577+
left, err := u.Left.TransformUp(f)
578+
if err != nil {
579+
return nil, err
580+
}
581+
582+
rigth, err := u.Right.TransformUp(f)
583+
if err != nil {
584+
return nil, err
585+
}
586+
587+
return f(NewUASTExtract(left, rigth))
588+
}

‎internal/function/uast_test.go

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,123 @@ func TestUASTXPath(t *testing.T) {
164164
}
165165
}
166166

167+
func TestUASTExtract(t *testing.T) {
168+
ctx, cleanup := setup(t)
169+
defer cleanup()
170+
171+
tests := []struct {
172+
name string
173+
key string
174+
expected []interface{}
175+
}{
176+
{
177+
name: "key_" + keyType,
178+
key: keyType,
179+
expected: []interface{}{
180+
[]string{"FunctionDef"},
181+
[]string{"Name"},
182+
[]string{"Name"},
183+
[]string{"Name"},
184+
[]string{"Name"},
185+
},
186+
},
187+
{
188+
name: "key_" + keyToken,
189+
key: keyToken,
190+
expected: []interface{}{
191+
[]string{"sum"},
192+
[]string{"a"},
193+
[]string{"b"},
194+
[]string{"print"},
195+
[]string{"sum"},
196+
},
197+
},
198+
{
199+
name: "key_" + keyRoles,
200+
key: keyRoles,
201+
expected: []interface{}{
202+
[]string{"Unannotated", "Function", "Declaration", "Name", "Identifier"},
203+
[]string{"Unannotated", "Identifier", "Expression", "Binary", "Left"},
204+
[]string{"Unannotated", "Identifier", "Expression", "Binary", "Right"},
205+
[]string{"Unannotated", "Identifier", "Expression", "Call", "Callee"},
206+
[]string{"Unannotated", "Identifier", "Expression", "Call", "Callee"},
207+
},
208+
},
209+
{
210+
name: "key_" + keyStartPos,
211+
key: keyStartPos,
212+
expected: []interface{}{
213+
[]string{"Offset:28 Line:4 Col:5 "},
214+
[]string{"Offset:47 Line:5 Col:9 "},
215+
[]string{"Offset:51 Line:5 Col:13 "},
216+
[]string{"Offset:54 Line:7 Col:1 "},
217+
[]string{"Offset:60 Line:7 Col:7 "},
218+
},
219+
},
220+
{
221+
name: "key_" + keyEndPos,
222+
key: keyEndPos,
223+
expected: []interface{}{
224+
[]string{"Offset:31 Line:4 Col:8 "},
225+
[]string{"Offset:48 Line:5 Col:10 "},
226+
[]string{"Offset:52 Line:5 Col:14 "},
227+
[]string{"Offset:59 Line:7 Col:6 "},
228+
[]string{"Offset:63 Line:7 Col:10 "},
229+
},
230+
},
231+
{
232+
name: "key_internalRole",
233+
key: "internalRole",
234+
expected: []interface{}{
235+
[]string{"body"},
236+
[]string{"left"},
237+
[]string{"right"},
238+
[]string{"func"},
239+
[]string{"func"},
240+
},
241+
},
242+
{
243+
name: "key_ctx",
244+
key: "ctx",
245+
expected: []interface{}{
246+
[]string{},
247+
[]string{"Load"},
248+
[]string{"Load"},
249+
[]string{"Load"},
250+
[]string{"Load"},
251+
},
252+
},
253+
{
254+
name: "key_foo",
255+
key: "foo",
256+
expected: []interface{}{
257+
[]string{},
258+
[]string{},
259+
[]string{},
260+
[]string{},
261+
[]string{},
262+
},
263+
},
264+
}
265+
266+
_, filteredNodes := bblfshFixtures(t, ctx)
267+
268+
for _, test := range tests {
269+
t.Run(test.name, func(t *testing.T) {
270+
row := sql.NewRow(filteredNodes["annotated"], test.key)
271+
272+
fn := NewUASTExtract(
273+
expression.NewGetField(0, sql.Array(sql.Blob), "", false),
274+
expression.NewLiteral(test.key, sql.Text),
275+
)
276+
277+
foo, err := fn.Eval(ctx, row)
278+
require.NoError(t, err)
279+
require.ElementsMatch(t, test.expected, foo)
280+
})
281+
}
282+
}
283+
167284
func assertUASTBlobs(t *testing.T, a, b interface{}) {
168285
t.Helper()
169286
require := require.New(t)

0 commit comments

Comments
 (0)