row2 changes

Author: James Cor

server/handler.go

@@ -495,6 +495,8 @@ func (h *Handler) doQuery(
 		r, err = resultForEmptyIter(sqlCtx, rowIter, resultFields)
 	} else if analyzer.FlagIsSet(qFlags, sql.QFlagMax1Row) {
 		r, err = resultForMax1RowIter(sqlCtx, schema, rowIter, resultFields, buf)
+	} else if ri2, ok := rowIter.(sql.RowIter2); ok && ri2.IsRowIter2(sqlCtx) {
+		r, processedAtLeastOneBatch, err = h.resultForDefaultIter2(sqlCtx, c, ri2, resultFields, callback, more)
 	} else {
 		r, processedAtLeastOneBatch, err = h.resultForDefaultIter(sqlCtx, c, schema, rowIter, callback, resultFields, more, buf)
 	}
@@ -768,6 +770,135 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 	return r, processedAtLeastOneBatch, nil
 }
 
+func (h *Handler) resultForDefaultIter2(ctx *sql.Context, c *mysql.Conn, iter sql.RowIter2, resultFields []*querypb.Field, callback func(*sqltypes.Result, bool) error, more bool) (*sqltypes.Result, bool, error) {
+	defer trace.StartRegion(ctx, "Handler.resultForDefaultIter2").End()
+
+	eg, ctx := ctx.NewErrgroup()
+	pan2err := func(err *error) {
+		if recoveredPanic := recover(); recoveredPanic != nil {
+			stack := debug.Stack()
+			wrappedErr := fmt.Errorf("handler caught panic: %v\n%s", recoveredPanic, stack)
+			*err = goerrors.Join(*err, wrappedErr)
+		}
+	}
+
+	// TODO: poll for closed connections should obviously also run even if
+	// we're doing something with an OK result or a single row result, etc.
+	// This should be in the caller.
+	pollCtx, cancelF := ctx.NewSubContext()
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
+		return h.pollForClosedConnection(pollCtx, c)
+	})
+
+	// Default waitTime is one minute if there is no timeout configured, in which case
+	// it will loop to iterate again unless the socket died by the OS timeout or other problems.
+	// If there is a timeout, it will be enforced to ensure that Vitess has a chance to
+	// call Handler.CloseConnection()
+	waitTime := 1 * time.Minute
+	if h.readTimeout > 0 {
+		waitTime = h.readTimeout
+	}
+	timer := time.NewTimer(waitTime)
+	defer timer.Stop()
+
+	wg := sync.WaitGroup{}
+	wg.Add(2)
+
+	// TODO: send results instead of rows?
+	// Read rows from iter and send them off
+	var rowChan = make(chan sql.Row2, 512)
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
+		defer wg.Done()
+		defer close(rowChan)
+		for {
+			select {
+			case <-ctx.Done():
+				return context.Cause(ctx)
+			default:
+				row, err := iter.Next2(ctx)
+				if err == io.EOF {
+					return nil
+				}
+				if err != nil {
+					return err
+				}
+				select {
+				case rowChan <- row:
+				case <-ctx.Done():
+					return nil
+				}
+			}
+		}
+	})
+
+	var res *sqltypes.Result
+	var processedAtLeastOneBatch bool
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
+		defer cancelF()
+		defer wg.Done()
+		for {
+			if res == nil {
+				res = &sqltypes.Result{
+					Fields: resultFields,
+					Rows:   make([][]sqltypes.Value, 0, rowsBatch),
+				}
+			}
+			if res.RowsAffected == rowsBatch {
+				if err := callback(res, more); err != nil {
+					return err
+				}
+				res = nil
+				processedAtLeastOneBatch = true
+				continue
+			}
+
+			select {
+			case <-ctx.Done():
+				return context.Cause(ctx)
+			case <-timer.C:
+				if h.readTimeout != 0 {
+					// Cancel and return so Vitess can call the CloseConnection callback
+					ctx.GetLogger().Tracef("connection timeout")
+					return ErrRowTimeout.New()
+				}
+			case row, ok := <-rowChan:
+				if !ok {
+					return nil
+				}
+				ctx.GetLogger().Tracef("spooling result row %s", row)
+				res.Rows = append(res.Rows, row)
+				res.RowsAffected++
+				if !timer.Stop() {
+					<-timer.C
+				}
+			}
+			timer.Reset(waitTime)
+		}
+	})
+
+	// Close() kills this PID in the process list,
+	// wait until all rows have be sent over the wire
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
+		wg.Wait()
+		return iter.Close(ctx)
+	})
+
+	err := eg.Wait()
+	if err != nil {
+		ctx.GetLogger().WithError(err).Warn("error running query")
+		if verboseErrorLogging {
+			fmt.Printf("Err: %+v", err)
+		}
+		return nil, false, err
+	}
+
+	return res, processedAtLeastOneBatch, nil
+}
+
 // See https://dev.mysql.com/doc/internals/en/status-flags.html
 func setConnStatusFlags(ctx *sql.Context, c *mysql.Conn) error {
 	ok, err := isSessionAutocommit(ctx)

sql/convert_value.go

@@ -3,98 +3,46 @@ package sql
 import (
 	"fmt"
 
-	"github.com/dolthub/vitess/go/vt/proto/query"
-
 	"github.com/dolthub/go-mysql-server/sql/values"
+
+	"github.com/dolthub/vitess/go/sqltypes"
+	"github.com/dolthub/vitess/go/vt/proto/query"
 )
 
 // ConvertToValue converts the interface to a sql value.
-func ConvertToValue(v interface{}) (Value, error) {
+func ConvertToValue(v interface{}) (sqltypes.Value, error) {
 	switch v := v.(type) {
 	case nil:
-		return Value{
-			Typ: query.Type_NULL_TYPE,
-			Val: nil,
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_NULL_TYPE, nil), nil
 	case int:
-		return Value{
-			Typ: query.Type_INT64,
-			Val: values.WriteInt64(make([]byte, values.Int64Size), int64(v)),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_INT64, values.WriteInt64(make([]byte, values.Int64Size), int64(v))), nil
 	case int8:
-		return Value{
-			Typ: query.Type_INT8,
-			Val: values.WriteInt8(make([]byte, values.Int8Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_INT8, values.WriteInt8(make([]byte, values.Int8Size), v)), nil
 	case int16:
-		return Value{
-			Typ: query.Type_INT16,
-			Val: values.WriteInt16(make([]byte, values.Int16Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_INT16, values.WriteInt16(make([]byte, values.Int16Size), v)), nil
 	case int32:
-		return Value{
-			Typ: query.Type_INT32,
-			Val: values.WriteInt32(make([]byte, values.Int32Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_INT32, values.WriteInt32(make([]byte, values.Int32Size), v)), nil
 	case int64:
-		return Value{
-			Typ: query.Type_INT64,
-			Val: values.WriteInt64(make([]byte, values.Int64Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_INT64, values.WriteInt64(make([]byte, values.Int64Size), v)), nil
 	case uint:
-		return Value{
-			Typ: query.Type_UINT64,
-			Val: values.WriteUint64(make([]byte, values.Uint64Size), uint64(v)),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_UINT64, values.WriteUint64(make([]byte, values.Uint64Size), uint64(v))), nil
 	case uint8:
-		return Value{
-			Typ: query.Type_UINT8,
-			Val: values.WriteUint8(make([]byte, values.Uint8Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_UINT8, values.WriteUint8(make([]byte, values.Uint8Size), v)), nil
 	case uint16:
-		return Value{
-			Typ: query.Type_UINT16,
-			Val: values.WriteUint16(make([]byte, values.Uint16Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_UINT16, values.WriteUint16(make([]byte, values.Uint16Size), v)), nil
 	case uint32:
-		return Value{
-			Typ: query.Type_UINT32,
-			Val: values.WriteUint32(make([]byte, values.Uint32Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_UINT32, values.WriteUint32(make([]byte, values.Uint32Size), v)), nil
 	case uint64:
-		return Value{
-			Typ: query.Type_UINT64,
-			Val: values.WriteUint64(make([]byte, values.Uint64Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_UINT64, values.WriteUint64(make([]byte, values.Uint64Size), v)), nil
 	case float32:
-		return Value{
-			Typ: query.Type_FLOAT32,
-			Val: values.WriteFloat32(make([]byte, values.Float32Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_FLOAT32, values.WriteFloat32(make([]byte, values.Float32Size), v)), nil
 	case float64:
-		return Value{
-			Typ: query.Type_FLOAT64,
-			Val: values.WriteFloat64(make([]byte, values.Float64Size), v),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_FLOAT64, values.WriteFloat64(make([]byte, values.Float64Size), v)), nil
 	case string:
-		return Value{
-			Typ: query.Type_VARCHAR,
-			Val: values.WriteString(make([]byte, len(v)), v, values.ByteOrderCollation),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_VARCHAR, values.WriteString(make([]byte, len(v)), v, values.ByteOrderCollation)), nil
 	case []byte:
-		return Value{
-			Typ: query.Type_BLOB,
-			Val: values.WriteBytes(make([]byte, len(v)), v, values.ByteOrderCollation),
-		}, nil
+		return sqltypes.MakeTrusted(query.Type_BLOB, values.WriteBytes(make([]byte, len(v)), v, values.ByteOrderCollation)), nil
 	default:
-		return Value{}, fmt.Errorf("type %T not implemented", v)
-	}
-}
-
-func MustConvertToValue(v interface{}) Value {
-	ret, err := ConvertToValue(v)
-	if err != nil {
-		panic(err)
+		return sqltypes.Value{}, fmt.Errorf("type %T not implemented", v)
 	}
-	return ret
 }

sql/core.go

@@ -26,6 +26,7 @@ import (
 	"time"
 	"unsafe"
 
+	"github.com/dolthub/vitess/go/sqltypes"
 	"github.com/shopspring/decimal"
 	"gopkg.in/src-d/go-errors.v1"
 
@@ -464,9 +465,10 @@ func DebugString(nodeOrExpression interface{}) string {
 type Expression2 interface {
 	Expression
 	// Eval2 evaluates the given row frame and returns a result.
-	Eval2(ctx *Context, row Row2) (Value, error)
+	Eval2(ctx *Context, row Row2) (sqltypes.Value, error)
 	// Type2 returns the expression type.
 	Type2() Type2
+	IsExpr2() bool
 }
 
 var SystemVariables SystemVariableRegistry

sql/expression/comparison.go

@@ -17,6 +17,8 @@ package expression
 import (
 	"fmt"
 
+	"github.com/dolthub/vitess/go/sqltypes"
+	querypb "github.com/dolthub/vitess/go/vt/proto/query"
 	errors "gopkg.in/src-d/go-errors.v1"
 
 	"github.com/dolthub/go-mysql-server/sql"
@@ -492,6 +494,7 @@ type GreaterThan struct {
 }
 
 var _ sql.Expression = (*GreaterThan)(nil)
+var _ sql.Expression2 = (*GreaterThan)(nil)
 var _ sql.CollationCoercible = (*GreaterThan)(nil)
 
 // NewGreaterThan creates a new GreaterThan expression.
@@ -518,6 +521,65 @@ func (gt *GreaterThan) Eval(ctx *sql.Context, row sql.Row) (interface{}, error)
 	return result == 1, nil
 }
 
+func (gt *GreaterThan) Eval2(ctx *sql.Context, row sql.Row2) (sqltypes.Value, error) {
+	l, ok := gt.Left().(sql.Expression2)
+	if !ok {
+		panic(fmt.Sprintf("%T does not implement sql.Expression2", gt.Left()))
+	}
+	r, ok := gt.Right().(sql.Expression2)
+	if !ok {
+		panic(fmt.Sprintf("%T does not implement sql.Expression2", gt.Right()))
+	}
+
+	lv, err := l.Eval2(ctx, row)
+	if err != nil {
+		return sqltypes.Value{}, err
+	}
+	rv, err := r.Eval2(ctx, row)
+	if err != nil {
+		return sqltypes.Value{}, err
+	}
+
+	// TODO: just assume they are int64
+	l64, err := types.ConvertValueToInt64(types.NumberTypeImpl_{}, lv)
+	if err != nil {
+		return sqltypes.Value{}, err
+	}
+	r64, err := types.ConvertValueToInt64(types.NumberTypeImpl_{}, rv)
+	if err != nil {
+		return sqltypes.Value{}, err
+	}
+	var rb byte
+	if l64 > r64 {
+		rb = 1
+	}
+
+	ret := sqltypes.MakeTrusted(querypb.Type_INT8, []byte{rb})
+	return ret, nil
+}
+
+func (gt *GreaterThan) Type2() sql.Type2 {
+	return nil
+}
+
+func (gt *GreaterThan) IsExpr2() bool {
+	lExpr, isExpr2 := gt.Left().(sql.Expression2)
+	if !isExpr2 {
+		return false
+	}
+	if !lExpr.IsExpr2() {
+		return false
+	}
+	rExpr, isExpr2 := gt.Right().(sql.Expression2)
+	if !isExpr2 {
+		return false
+	}
+	if !rExpr.IsExpr2() {
+		return false
+	}
+	return true
+}
+
 // WithChildren implements the Expression interface.
 func (gt *GreaterThan) WithChildren(children ...sql.Expression) (sql.Expression, error) {
 	if len(children) != 2 {

sql/expression/get_field.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/dolthub/vitess/go/sqltypes"
 	errors "gopkg.in/src-d/go-errors.v1"
 
 	"github.com/dolthub/go-mysql-server/sql"
@@ -149,12 +150,15 @@ func (p *GetField) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
 	return row[p.fieldIndex], nil
 }
 
-func (p *GetField) Eval2(ctx *sql.Context, row sql.Row2) (sql.Value, error) {
+func (p *GetField) Eval2(ctx *sql.Context, row sql.Row2) (sqltypes.Value, error) {
 	if p.fieldIndex < 0 || p.fieldIndex >= row.Len() {
-		return sql.Value{}, ErrIndexOutOfBounds.New(p.fieldIndex, row.Len())
+		return sqltypes.Value{}, ErrIndexOutOfBounds.New(p.fieldIndex, row.Len())
 	}
+	return row[p.fieldIndex], nil
+}
 
-	return row.GetField(p.fieldIndex), nil
+func (p *GetField) IsExpr2() bool {
+	return true
 }
 
 // WithChildren implements the Expression interface.