function: LOC function implementation

Signed-off-by: Máximo Cuadros <[email protected]>

internal/function/loc.go

@@ -0,0 +1,156 @@
+package function
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+
+	"github.com/hhatto/gocloc"
+	"gopkg.in/src-d/enry.v1"
+	"gopkg.in/src-d/go-mysql-server.v0/sql"
+)
+
+var languages = gocloc.NewDefinedLanguages()
+
+var errEmptyInputValues = errors.New("empty input values")
+
+type LOC struct {
+	Left  sql.Expression
+	Right sql.Expression
+}
+
+// NewLOC creates a new LOC UDF.
+func NewLOC(args ...sql.Expression) (sql.Expression, error) {
+	if len(args) != 2 {
+		return nil, sql.ErrInvalidArgumentNumber.New("2", len(args))
+	}
+
+	return &LOC{args[0], args[1]}, nil
+}
+
+// Resolved implements the Expression interface.
+func (f *LOC) Resolved() bool {
+	return f.Left.Resolved() && f.Right.Resolved()
+}
+
+func (f *LOC) String() string {
+	return fmt.Sprintf("loc(%s, %s)", f.Left, f.Right)
+}
+
+// IsNullable implements the Expression interface.
+func (f *LOC) IsNullable() bool {
+	return f.Left.IsNullable() || f.Right.IsNullable()
+}
+
+// Type implements the Expression interface.
+func (LOC) Type() sql.Type {
+	return sql.JSON
+}
+
+// TransformUp implements the Expression interface.
+func (f *LOC) TransformUp(fn sql.TransformExprFunc) (sql.Expression, error) {
+	left, err := f.Left.TransformUp(fn)
+	if err != nil {
+		return nil, err
+	}
+
+	right, err := f.Right.TransformUp(fn)
+	if err != nil {
+		return nil, err
+	}
+
+	return fn(&LOC{left, right})
+}
+
+// Eval implements the Expression interface.
+func (f *LOC) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	span, ctx := ctx.Span("gitbase.LOC")
+	defer span.Finish()
+	path, blob, err := f.getInputValues(ctx, row)
+	if err != nil {
+		if err == errEmptyInputValues {
+			return nil, nil
+		}
+
+		return nil, err
+	}
+
+	lang := f.getLanguage(path, blob)
+	if lang == "" || languages.Langs[lang] == nil {
+		return nil, nil
+	}
+
+	return gocloc.AnalyzeReader(
+		path,
+		languages.Langs[lang],
+		bytes.NewReader(blob), &gocloc.ClocOptions{},
+	), nil
+}
+
+func (f *LOC) getInputValues(ctx *sql.Context, row sql.Row) (string, []byte, error) {
+	left, err := f.Left.Eval(ctx, row)
+	if err != nil {
+		return "", nil, err
+	}
+
+	left, err = sql.Text.Convert(left)
+	if err != nil {
+		return "", nil, err
+	}
+
+	right, err := f.Right.Eval(ctx, row)
+	if err != nil {
+		return "", nil, err
+	}
+
+	right, err = sql.Blob.Convert(right)
+	if err != nil {
+		return "", nil, err
+	}
+
+	if right == nil {
+		return "", nil, errEmptyInputValues
+	}
+
+	path, ok := left.(string)
+	if !ok {
+		return "", nil, errEmptyInputValues
+	}
+
+	blob, ok := right.([]byte)
+
+	if !ok {
+		return "", nil, errEmptyInputValues
+	}
+
+	if len(blob) == 0 || len(path) == 0 {
+		return "", nil, errEmptyInputValues
+	}
+
+	return path, blob, nil
+}
+
+func (f *LOC) getLanguage(path string, blob []byte) string {
+	hash := languageHash(path, blob)
+
+	value, ok := languageCache.Get(hash)
+	if ok {
+		return value.(string)
+	}
+
+	lang := enry.GetLanguage(path, blob)
+	if len(blob) > 0 {
+		languageCache.Add(hash, lang)
+	}
+
+	return lang
+}
+
+// Children implements the Expression interface.
+func (f *LOC) Children() []sql.Expression {
+	if f.Right == nil {
+		return []sql.Expression{f.Left}
+	}
+
+	return []sql.Expression{f.Left, f.Right}
+}

internal/function/loc_test.go

@@ -0,0 +1,56 @@
+package function
+
+import (
+	"testing"
+
+	"github.com/hhatto/gocloc"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/src-d/go-errors.v1"
+	"gopkg.in/src-d/go-mysql-server.v0/sql"
+	"gopkg.in/src-d/go-mysql-server.v0/sql/expression"
+)
+
+func TestLoc(t *testing.T) {
+	testCases := []struct {
+		name     string
+		row      sql.Row
+		expected interface{}
+		err      *errors.Kind
+	}{
+		{"left is null", sql.NewRow(nil), nil, nil},
+		{"both are null", sql.NewRow(nil, nil), nil, nil},
+		{"too few args given", sql.NewRow("foo.foobar"), nil, nil},
+		{"too many args given", sql.NewRow("foo.rb", "bar", "baz"), nil, sql.ErrInvalidArgumentNumber},
+		{"invalid blob type given", sql.NewRow("foo", 5), nil, sql.ErrInvalidType},
+		{"path and blob are given", sql.NewRow("foo", "#!/usr/bin/env python\n\nprint 'foo'"), &gocloc.ClocFile{
+			Code: 2, Comments: 0, Blanks: 1, Name: "foo", Lang: "",
+		}, nil},
+	}
+
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+			ctx := sql.NewEmptyContext()
+
+			var args = make([]sql.Expression, len(tt.row))
+			for i := range tt.row {
+				args[i] = expression.NewGetField(i, sql.Text, "", false)
+			}
+
+			f, err := NewLOC(args...)
+			if err == nil {
+				var val interface{}
+				val, err = f.Eval(ctx, tt.row)
+				if tt.err == nil {
+					require.NoError(err)
+					require.Equal(tt.expected, val)
+				}
+			}
+
+			if tt.err != nil {
+				require.Error(err)
+				require.True(tt.err.Is(err))
+			}
+		})
+	}
+}

internal/function/registry.go

@@ -7,6 +7,7 @@ var Functions = []sql.Function{
 	sql.Function1{Name: "is_tag", Fn: NewIsTag},
 	sql.Function1{Name: "is_remote", Fn: NewIsRemote},
 	sql.FunctionN{Name: "language", Fn: NewLanguage},
+	sql.FunctionN{Name: "loc", Fn: NewLOC},
 	sql.FunctionN{Name: "uast", Fn: NewUAST},
 	sql.Function3{Name: "uast_mode", Fn: NewUASTMode},
 	sql.Function2{Name: "uast_xpath", Fn: NewUASTXPath},