src-d · ajnavarro · Nov 8, 2018 · Nov 8, 2018
@@ -58,4 +58,7 @@ var Defaults = sql.Functions{
 	"sqrt":          sql.Function1(NewSqrt),
 	"pow":           sql.Function2(NewPower),
 	"power":         sql.Function2(NewPower),
+	"ltrim":         sql.Function1(NewTrimFunc(lTrimType)),
+	"rtrim":         sql.Function1(NewTrimFunc(rTrimType)),
+	"trim":          sql.Function1(NewTrimFunc(bTrimType)),
 }
@@ -0,0 +1,94 @@
+package function
+
+import (
+	"fmt"
+	"reflect"
+	"strings"
+	"unicode"
+
+	"gopkg.in/src-d/go-mysql-server.v0/sql/expression"
+	"gopkg.in/src-d/go-mysql-server.v0/sql"
+)
+
+type trimType rune
+const (
+	lTrimType trimType = 'l'
+	rTrimType trimType = 'r'
+	bTrimType trimType = 'b'
+)
+
+// NewTrimFunc returns a Trim creator function with a specific trimType.
+func NewTrimFunc(tType trimType) func(e sql.Expression) sql.Expression {
+	return func(e sql.Expression) sql.Expression {
+		return NewTrim(tType, e)
+	}
+}
+
+// NewTrim creates a new Trim expression.
+func NewTrim(tType trimType, str sql.Expression) sql.Expression {
+	return &Trim{expression.UnaryExpression{Child: str}, tType}
+}
+
+// Trim is a function that returns the string with prefix or suffix spaces removed based on the trimType
+type Trim struct {
+	expression.UnaryExpression
+	trimType
+}
+
+// Type implements the Expression interface.
+func (t *Trim) Type() sql.Type { return sql.Text }
+
+func (t *Trim) String() string {
+	switch t.trimType {
+	case lTrimType:
+		return fmt.Sprintf("ltrim(%s)", t.Child)
+	case rTrimType:
+		return fmt.Sprintf("rtrim(%s)", t.Child)
+	default:
+		return fmt.Sprintf("trim(%s)", t.Child)
+	}
+}
+
+// IsNullable implements the Expression interface.
+func (t *Trim) IsNullable() bool {
+	return t.Child.IsNullable()
+}
+
+// TransformUp implements the Expression interface.
+func (t *Trim) TransformUp(f sql.TransformExprFunc) (sql.Expression, error) {
+	str, err := t.Child.TransformUp(f)
+	if err != nil {
+		return nil, err
+	}
+
+	return f(NewTrim(t.trimType, str))
+}
+
+// Eval implements the Expression interface.
+func (t *Trim) Eval(
+	ctx *sql.Context,
+	row sql.Row,
+) (interface{}, error) {
+	str, err := t.Child.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+
+	if str == nil {
+		return nil, nil
+	}
+
+	str, err = sql.Text.Convert(str)
+	if err != nil {
+		return nil, sql.ErrInvalidType.New(reflect.TypeOf(str))
+	}
+
+	switch t.trimType {
+	case lTrimType:
+		return strings.TrimLeftFunc(str.(string), unicode.IsSpace), nil
+	case rTrimType:
+		return strings.TrimRightFunc(str.(string), unicode.IsSpace), nil
+	default:
+		return strings.TrimFunc(str.(string), unicode.IsSpace), nil
+	}
+}
@@ -0,0 +1,108 @@
+package function
+
+import (
+	"testing"
+
+	"gopkg.in/src-d/go-mysql-server.v0/sql/expression"
+	"gopkg.in/src-d/go-mysql-server.v0/sql"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTrim(t *testing.T) {
+	f := NewTrimFunc(bTrimType)(expression.NewGetField(0, sql.Text, "", false), )
+	testCases := []struct {
+		name     string
+		row      sql.Row
+		expected interface{}
+		err      bool
+	}{
+		{"null input", sql.NewRow(nil), nil, false},
+		{"trimmed string", sql.NewRow("foo"), "foo", false},
+		{"spaces in both sides", sql.NewRow("  foo    "), "foo", false},
+		{"spaces in left side", sql.NewRow("  foo"), "foo", false},
+		{"spaces in right side", sql.NewRow("foo    "), "foo", false},
+		{"two words with spaces", sql.NewRow(" foo   bar "), "foo   bar", false},
+		{"different kinds of spaces", sql.NewRow("\r\tfoo   bar \v"), "foo   bar", false},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Helper()
+			require := require.New(t)
+			ctx := sql.NewEmptyContext()
+
+			v, err := f.Eval(ctx, tt.row)
+			if tt.err {
+				require.Error(err)
+			} else {
+				require.NoError(err)
+				require.Equal(tt.expected, v)
+			}
+		})
+	}
+}
+
+func TestLTrim(t *testing.T) {
+	f := NewTrimFunc(lTrimType)(expression.NewGetField(0, sql.Text, "", false), )
+	testCases := []struct {
+		name     string
+		row      sql.Row
+		expected interface{}
+		err      bool
+	}{
+		{"null input", sql.NewRow(nil), nil, false},
+		{"trimmed string", sql.NewRow("foo"), "foo", false},
+		{"spaces in both sides", sql.NewRow("  foo    "), "foo    ", false},
+		{"spaces in left side", sql.NewRow("  foo"), "foo", false},
+		{"spaces in right side", sql.NewRow("foo    "), "foo    ", false},
+		{"two words with spaces", sql.NewRow(" foo   bar "), "foo   bar ", false},
+		{"different kinds of spaces", sql.NewRow("\r\tfoo   bar \v"), "foo   bar \v", false},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Helper()
+			require := require.New(t)
+			ctx := sql.NewEmptyContext()
+
+			v, err := f.Eval(ctx, tt.row)
+			if tt.err {
+				require.Error(err)
+			} else {
+				require.NoError(err)
+				require.Equal(tt.expected, v)
+			}
+		})
+	}
+}
+
+func TestRTrim(t *testing.T) {
+	f := NewTrimFunc(rTrimType)(expression.NewGetField(0, sql.Text, "", false), )
+	testCases := []struct {
+		name     string
+		row      sql.Row
+		expected interface{}
+		err      bool
+	}{
+		{"null input", sql.NewRow(nil), nil, false},
+		{"trimmed string", sql.NewRow("foo"), "foo", false},
+		{"spaces in both sides", sql.NewRow("  foo    "), "  foo", false},
+		{"spaces in left side", sql.NewRow("  foo"), "  foo", false},
+		{"spaces in right side", sql.NewRow("foo    "), "foo", false},
+		{"two words with spaces", sql.NewRow(" foo   bar "), " foo   bar", false},
+		{"different kinds of spaces", sql.NewRow("\r\tfoo   bar \v"), "\r\tfoo   bar", false},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Helper()
+			require := require.New(t)
+			ctx := sql.NewEmptyContext()
+
+			v, err := f.Eval(ctx, tt.row)
+			if tt.err {
+				require.Error(err)
+			} else {
+				require.NoError(err)
+				require.Equal(tt.expected, v)
+			}
+		})
+	}
+}