Merge pull request #446 from mcarmonaa/feature/udf-uast-extract

internal/function: add uast_extract udf

Author: ajnavarro

docs/using-gitbase/functions.md

@@ -12,6 +12,7 @@ To make some common tasks easier for the user, there are some functions to inter
 |uast(blob, [lang, [xpath]])json_blob| returns an array of UAST nodes as blobs in semantic mode                                                 |
 |uast_mode(mode, blob, lang)json_blob| returns an array of UAST nodes as blobs specifying its language and mode (semantic, annotated or native) |
 |uast_xpath(json_blob, xpath)| performs an XPath query over the given UAST nodes                                                                |
+|uast_extract(json_blob, key)| extracts information identified by the given key from the uast nodes                                              |
 
 ## Standard functions
 

internal/function/registry.go

@@ -4,10 +4,11 @@ import "gopkg.in/src-d/go-mysql-server.v0/sql"
 
 // Functions for gitbase queries.
 var Functions = sql.Functions{
-	"is_tag":     sql.Function1(NewIsTag),
-	"is_remote":  sql.Function1(NewIsRemote),
-	"language":   sql.FunctionN(NewLanguage),
-	"uast":       sql.FunctionN(NewUAST),
-	"uast_mode":  sql.Function3(NewUASTMode),
-	"uast_xpath": sql.Function2(NewUASTXPath),
+	"is_tag":       sql.Function1(NewIsTag),
+	"is_remote":    sql.Function1(NewIsRemote),
+	"language":     sql.FunctionN(NewLanguage),
+	"uast":         sql.FunctionN(NewUAST),
+	"uast_mode":    sql.Function3(NewUASTMode),
+	"uast_xpath":   sql.Function2(NewUASTXPath),
+	"uast_extract": sql.Function2(NewUASTExtract),
 }

internal/function/uast.go

@@ -275,12 +275,7 @@ func (f UASTMode) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err erro
 		return nil, fmt.Errorf("invalid uast mode %s", m)
 	}
 
-	u, err := getUAST(ctx, bytes, lang, "", mode)
-	if err != nil {
-		return nil, err
-	}
-
-	return u, nil
+	return getUAST(ctx, bytes, lang, "", mode)
 }
 
 // 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
 		return nil, nil
 	}
 
-	left, err = sql.Array(sql.Blob).Convert(left)
+	nodes, err := nodesFromBlobArray(left)
 	if err != nil {
 		return nil, err
 	}
 
-	arr := left.([]interface{})
-	var nodes = make([]*uast.Node, len(arr))
-	for i, n := range arr {
-		node := uast.NewNode()
-		if err := node.Unmarshal(n.([]byte)); err != nil {
-			return nil, err
-		}
-		nodes[i] = node
-	}
-
-	right, err := f.Right.Eval(ctx, row)
+	xpath, err := exprToString(ctx, f.Right, row)
 	if err != nil {
 		return nil, err
 	}
 
-	if right == nil {
+	if xpath == "" {
 		return nil, nil
 	}
 
-	right, err = sql.Text.Convert(right)
-	if err != nil {
-		return nil, err
-	}
-
-	xpath := right.(string)
-
 	var result []interface{}
 	for _, n := range nodes {
 		ns, err := tools.Filter(n, xpath)
@@ -368,6 +346,25 @@ func (f *UASTXPath) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err er
 	return result, nil
 }
 
+func nodesFromBlobArray(data interface{}) ([]*uast.Node, error) {
+	data, err := sql.Array(sql.Blob).Convert(data)
+	if err != nil {
+		return nil, err
+	}
+
+	arr := data.([]interface{})
+	var nodes = make([]*uast.Node, len(arr))
+	for i, n := range arr {
+		node := uast.NewNode()
+		if err := node.Unmarshal(n.([]byte)); err != nil {
+			return nil, err
+		}
+		nodes[i] = node
+	}
+
+	return nodes, nil
+}
+
 func (f UASTXPath) String() string {
 	return fmt.Sprintf("uast_xpath(%s, %s)", f.Left, f.Right)
 }
@@ -475,3 +472,117 @@ func getUAST(
 
 	return result, nil
 }
+
+// UASTExtract extracts keys from an UAST.
+type UASTExtract struct {
+	expression.BinaryExpression
+}
+
+// NewUASTExtract creates a new UASTExtract UDF.
+func NewUASTExtract(uast, key sql.Expression) sql.Expression {
+	return &UASTExtract{expression.BinaryExpression{Left: uast, Right: key}}
+}
+
+// String implements the fmt.Stringer interface.
+func (u *UASTExtract) String() string {
+	return fmt.Sprintf("uast_extract(%s, %s)", u.Left, u.Right)
+}
+
+// Type implements the sql.Expression interface.
+func (u *UASTExtract) Type() sql.Type {
+	return sql.Array(sql.Array(sql.Text))
+}
+
+// Eval implements the sql.Expression interface.
+func (u *UASTExtract) Eval(ctx *sql.Context, row sql.Row) (out interface{}, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			err = fmt.Errorf("uast: unknown error: %s", r)
+		}
+	}()
+
+	span, ctx := ctx.Span("gitbase.UASTExtract")
+	defer span.Finish()
+
+	left, err := u.Left.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+
+	if left == nil {
+		return nil, nil
+	}
+
+	nodes, err := nodesFromBlobArray(left)
+	if err != nil {
+		return nil, err
+	}
+
+	key, err := exprToString(ctx, u.Right, row)
+	if err != nil {
+		return nil, err
+	}
+
+	if key == "" {
+		return nil, nil
+	}
+
+	extracted := make([][]string, len(nodes))
+	for i, n := range nodes {
+		extracted[i] = extractInfo(n, key)
+	}
+
+	return extracted, nil
+}
+
+const (
+	keyType     = "@type"
+	keyToken    = "@token"
+	keyRoles    = "@role"
+	keyStartPos = "@startpos"
+	keyEndPos   = "@endpos"
+)
+
+func extractInfo(n *uast.Node, key string) []string {
+
+	info := []string{}
+	switch key {
+	case keyType:
+		info = append(info, n.InternalType)
+	case keyToken:
+		info = append(info, n.Token)
+	case keyRoles:
+		roles := make([]string, len(n.Roles))
+		for i, rol := range n.Roles {
+			roles[i] = rol.String()
+		}
+
+		info = append(info, roles...)
+	case keyStartPos:
+		info = append(info, n.StartPosition.String())
+	case keyEndPos:
+		info = append(info, n.EndPosition.String())
+	default:
+		v, ok := n.Properties[key]
+		if ok {
+			info = append(info, v)
+		}
+	}
+
+	return info
+}
+
+// TransformUp implements the sql.Expression interface.
+func (u *UASTExtract) TransformUp(f sql.TransformExprFunc) (sql.Expression, error) {
+	left, err := u.Left.TransformUp(f)
+	if err != nil {
+		return nil, err
+	}
+
+	rigth, err := u.Right.TransformUp(f)
+	if err != nil {
+		return nil, err
+	}
+
+	return f(NewUASTExtract(left, rigth))
+}

internal/function/uast_test.go

@@ -164,6 +164,123 @@ func TestUASTXPath(t *testing.T) {
 	}
 }
 
+func TestUASTExtract(t *testing.T) {
+	ctx, cleanup := setup(t)
+	defer cleanup()
+
+	tests := []struct {
+		name     string
+		key      string
+		expected []interface{}
+	}{
+		{
+			name: "key_" + keyType,
+			key:  keyType,
+			expected: []interface{}{
+				[]string{"FunctionDef"},
+				[]string{"Name"},
+				[]string{"Name"},
+				[]string{"Name"},
+				[]string{"Name"},
+			},
+		},
+		{
+			name: "key_" + keyToken,
+			key:  keyToken,
+			expected: []interface{}{
+				[]string{"sum"},
+				[]string{"a"},
+				[]string{"b"},
+				[]string{"print"},
+				[]string{"sum"},
+			},
+		},
+		{
+			name: "key_" + keyRoles,
+			key:  keyRoles,
+			expected: []interface{}{
+				[]string{"Unannotated", "Function", "Declaration", "Name", "Identifier"},
+				[]string{"Unannotated", "Identifier", "Expression", "Binary", "Left"},
+				[]string{"Unannotated", "Identifier", "Expression", "Binary", "Right"},
+				[]string{"Unannotated", "Identifier", "Expression", "Call", "Callee"},
+				[]string{"Unannotated", "Identifier", "Expression", "Call", "Callee"},
+			},
+		},
+		{
+			name: "key_" + keyStartPos,
+			key:  keyStartPos,
+			expected: []interface{}{
+				[]string{"Offset:28 Line:4 Col:5 "},
+				[]string{"Offset:47 Line:5 Col:9 "},
+				[]string{"Offset:51 Line:5 Col:13 "},
+				[]string{"Offset:54 Line:7 Col:1 "},
+				[]string{"Offset:60 Line:7 Col:7 "},
+			},
+		},
+		{
+			name: "key_" + keyEndPos,
+			key:  keyEndPos,
+			expected: []interface{}{
+				[]string{"Offset:31 Line:4 Col:8 "},
+				[]string{"Offset:48 Line:5 Col:10 "},
+				[]string{"Offset:52 Line:5 Col:14 "},
+				[]string{"Offset:59 Line:7 Col:6 "},
+				[]string{"Offset:63 Line:7 Col:10 "},
+			},
+		},
+		{
+			name: "key_internalRole",
+			key:  "internalRole",
+			expected: []interface{}{
+				[]string{"body"},
+				[]string{"left"},
+				[]string{"right"},
+				[]string{"func"},
+				[]string{"func"},
+			},
+		},
+		{
+			name: "key_ctx",
+			key:  "ctx",
+			expected: []interface{}{
+				[]string{},
+				[]string{"Load"},
+				[]string{"Load"},
+				[]string{"Load"},
+				[]string{"Load"},
+			},
+		},
+		{
+			name: "key_foo",
+			key:  "foo",
+			expected: []interface{}{
+				[]string{},
+				[]string{},
+				[]string{},
+				[]string{},
+				[]string{},
+			},
+		},
+	}
+
+	_, filteredNodes := bblfshFixtures(t, ctx)
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			row := sql.NewRow(filteredNodes["annotated"], test.key)
+
+			fn := NewUASTExtract(
+				expression.NewGetField(0, sql.Array(sql.Blob), "", false),
+				expression.NewLiteral(test.key, sql.Text),
+			)
+
+			foo, err := fn.Eval(ctx, row)
+			require.NoError(t, err)
+			require.ElementsMatch(t, test.expected, foo)
+		})
+	}
+}
+
 func assertUASTBlobs(t *testing.T, a, b interface{}) {
 	t.Helper()
 	require := require.New(t)