Implement sql.ValueRow (#3248)

Author: jycor

go.mod

@@ -14,7 +14,7 @@ require (
 	github.com/lestrrat-go/strftime v1.0.4
 	github.com/pkg/errors v0.9.1
 	github.com/pmezard/go-difflib v1.0.0
-	github.com/shopspring/decimal v1.3.1
+	github.com/shopspring/decimal v1.4.0
 	github.com/sirupsen/logrus v1.8.1
 	github.com/stretchr/testify v1.9.0
 	go.opentelemetry.io/otel v1.31.0

go.sum

@@ -66,8 +66,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/shopspring/decimal v1.3.1 h1:2Usl1nmF/WZucqkFZhnfFYxxxu8LG21F6nPQBE5gKV8=
-github.com/shopspring/decimal v1.3.1/go.mod h1:DKyhrW/HYNuLGql+MJL6WCR6knT2jwCFRcu2hWCYk4o=
+github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k=
+github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME=
 github.com/sirupsen/logrus v1.8.1 h1:dJKuHgqk1NNQlqoA6BTlM1Wf9DOH3NBjQyu0h9+AZZE=
 github.com/sirupsen/logrus v1.8.1/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=

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 vr, ok := rowIter.(sql.ValueRowIter); ok && vr.IsValueRowIter(sqlCtx) {
+		r, processedAtLeastOneBatch, err = h.resultForValueRowIter(sqlCtx, c, schema, vr, resultFields, buf, callback, more)
 	} else {
 		r, processedAtLeastOneBatch, err = h.resultForDefaultIter(sqlCtx, c, schema, rowIter, callback, resultFields, more, buf)
 	}
@@ -768,6 +770,149 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 	return r, processedAtLeastOneBatch, nil
 }
 
+func (h *Handler) resultForValueRowIter(ctx *sql.Context, c *mysql.Conn, schema sql.Schema, iter sql.ValueRowIter, resultFields []*querypb.Field, buf *sql.ByteBuffer, callback func(*sqltypes.Result, bool) error, more bool) (*sqltypes.Result, bool, error) {
+	defer trace.StartRegion(ctx, "Handler.resultForValueRowIter").End()
+
+	eg, ctx := ctx.NewErrgroup()
+	pan2err := func(err *error) {
+		if recoveredPanic := recover(); recoveredPanic != nil {
+			wrappedErr := fmt.Errorf("handler caught panic: %v\n%s", recoveredPanic, debug.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)
+
+	// Wrap the callback to include a BytesBuffer.Reset() for non-cursor requests, to
+	// clean out rows that have already been spooled.
+	resetCallback := func(r *sqltypes.Result, more bool) error {
+		// A server-side cursor allows the caller to fetch results cached on the server-side,
+		// so if a cursor exists, we can't release the buffer memory yet.
+		if c.StatusFlags&uint16(mysql.ServerCursorExists) != 0 {
+			defer buf.Reset()
+		}
+		return callback(r, more)
+	}
+
+	// TODO: send results instead of rows?
+	// Read rows from iter and send them off
+	var rowChan = make(chan sql.ValueRow, 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.NextValueRow(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 := resetCallback(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
+				}
+				resRow, err := RowValueToSQLValues(ctx, schema, row, buf)
+				if err != nil {
+					return err
+				}
+				ctx.GetLogger().Tracef("spooling result row %s", resRow)
+				res.Rows = append(res.Rows, resRow)
+				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)
@@ -994,7 +1139,7 @@ func toSqlHelper(ctx *sql.Context, typ sql.Type, buf *sql.ByteBuffer, val interf
 		return typ.SQL(ctx, nil, val)
 	}
 	ret, err := typ.SQL(ctx, buf.Get(), val)
-	buf.Grow(ret.Len())
+	buf.Grow(ret.Len()) // TODO: shouldn't we check capacity beforehand?
 	return ret, err
 }
 
@@ -1037,6 +1182,39 @@ func RowToSQL(ctx *sql.Context, sch sql.Schema, row sql.Row, projs []sql.Express
 	return outVals, nil
 }
 
+func RowValueToSQLValues(ctx *sql.Context, sch sql.Schema, row sql.ValueRow, buf *sql.ByteBuffer) ([]sqltypes.Value, error) {
+	if len(sch) == 0 {
+		return []sqltypes.Value{}, nil
+	}
+	var err error
+	outVals := make([]sqltypes.Value, len(sch))
+	for i, col := range sch {
+		// TODO: remove this check once all Types implement this
+		valType, ok := col.Type.(sql.ValueType)
+		if !ok {
+			if row[i].IsNull() {
+				outVals[i] = sqltypes.NULL
+				continue
+			}
+			outVals[i] = sqltypes.MakeTrusted(row[i].Typ, row[i].Val)
+			continue
+		}
+		if buf == nil {
+			outVals[i], err = valType.SQLValue(ctx, row[i], nil)
+			if err != nil {
+				return nil, err
+			}
+			continue
+		}
+		outVals[i], err = valType.SQLValue(ctx, row[i], buf.Get())
+		if err != nil {
+			return nil, err
+		}
+		buf.Grow(outVals[i].Len())
+	}
+	return outVals, nil
+}
+
 func schemaToFields(ctx *sql.Context, s sql.Schema) []*querypb.Field {
 	charSetResults := ctx.GetCharacterSetResults()
 	fields := make([]*querypb.Field, len(s))

sql/analyzer/replace_sort.go

@@ -175,12 +175,12 @@ func replaceIdxSortHelper(ctx *sql.Context, scope *plan.Scope, node sql.Node, so
 					sortFields[i] = sortField
 				} else {
 					sameSortFields = false
-					col2, _ := col.(sql.Expression2)
+					valCol, _ := col.(sql.ValueExpression)
 					sortFields[i] = sql.SortField{
-						Column:       col,
-						Column2:      col2,
-						NullOrdering: sortField.NullOrdering,
-						Order:        sortField.Order,
+						Column:          col,
+						ValueExprColumn: valCol,
+						NullOrdering:    sortField.NullOrdering,
+						Order:           sortField.Order,
 					}
 				}
 			}

sql/cache.go

@@ -71,10 +71,10 @@ func (l *lruCache) Dispose() {
 }
 
 type rowsCache struct {
-	memory   Freeable
-	reporter Reporter
-	rows     []Row
-	rows2    []Row2
+	memory    Freeable
+	reporter  Reporter
+	rows      []Row
+	valueRows []ValueRow
 }
 
 func newRowsCache(memory Freeable, r Reporter) *rowsCache {
@@ -92,17 +92,17 @@ func (c *rowsCache) Add(row Row) error {
 
 func (c *rowsCache) Get() []Row { return c.rows }
 
-func (c *rowsCache) Add2(row2 Row2) error {
+func (c *rowsCache) AddValueRow(row ValueRow) error {
 	if !releaseMemoryIfNeeded(c.reporter, c.memory.Free) {
 		return ErrNoMemoryAvailable.New()
 	}
 
-	c.rows2 = append(c.rows2, row2)
+	c.valueRows = append(c.valueRows, row)
 	return nil
 }
 
-func (c *rowsCache) Get2() []Row2 {
-	return c.rows2
+func (c *rowsCache) GetValueRow() []ValueRow {
+	return c.valueRows
 }
 
 func (c *rowsCache) Dispose() {

sql/convert_value.go

@@ -3,9 +3,9 @@ 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/vt/proto/query"
 )
 
 // ConvertToValue converts the interface to a sql value.
@@ -90,11 +90,3 @@ func ConvertToValue(v interface{}) (Value, error) {
 		return Value{}, fmt.Errorf("type %T not implemented", v)
 	}
 }
-
-func MustConvertToValue(v interface{}) Value {
-	ret, err := ConvertToValue(v)
-	if err != nil {
-		panic(err)
-	}
-	return ret
-}

sql/core.go

@@ -460,13 +460,13 @@ func DebugString(nodeOrExpression interface{}) string {
 	panic(fmt.Sprintf("Expected sql.DebugString or fmt.Stringer for %T", nodeOrExpression))
 }
 
-// Expression2 is an experimental future interface alternative to Expression to provide faster access.
-type Expression2 interface {
+// ValueExpression is an experimental future interface alternative to Expression to provide faster access.
+type ValueExpression interface {
 	Expression
-	// Eval2 evaluates the given row frame and returns a result.
-	Eval2(ctx *Context, row Row2) (Value, error)
-	// Type2 returns the expression type.
-	Type2() Type2
+	// EvalValue evaluates the given row frame and returns a result.
+	EvalValue(ctx *Context, row ValueRow) (Value, error)
+	// IsValueExpression indicates whether this expression and all its children support ValueExpression.
+	IsValueExpression() bool
 }
 
 var SystemVariables SystemVariableRegistry