implement row2

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, err = h.resultForDefaultIter2(sqlCtx, ri2, resultFields, callback, more)
 	} else {
 		r, processedAtLeastOneBatch, err = h.resultForDefaultIter(sqlCtx, c, schema, rowIter, callback, resultFields, more, buf)
 	}
@@ -768,6 +770,32 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 	return r, processedAtLeastOneBatch, nil
 }
 
+func (h *Handler) resultForDefaultIter2(ctx *sql.Context, iter sql.RowIter2, resultFields []*querypb.Field, callback func(*sqltypes.Result, bool) error, more bool) (*sqltypes.Result, error) {
+	res := &sqltypes.Result{Fields: resultFields}
+	for {
+		if res.RowsAffected == rowsBatch {
+			if err := callback(res, more); err != nil {
+				return nil, err
+			}
+			res = nil
+		}
+		row, err := iter.Next2(ctx)
+		if err == io.EOF {
+			return res, nil
+		}
+		if err != nil {
+			return nil, err
+		}
+
+		outRow := make([]sqltypes.Value, len(res.Rows))
+		for i := range row {
+			outRow[i] = sqltypes.MakeTrusted(row[i].Typ, row[i].Val)
+		}
+		res.Rows = append(res.Rows, outRow)
+		res.RowsAffected++
+	}
+}
+
 // 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/plan/process.go

@@ -317,6 +317,29 @@ func (i *TrackedRowIter) Next(ctx *sql.Context) (sql.Row, error) {
 	return row, nil
 }
 
+func (i *TrackedRowIter) Next2(ctx *sql.Context) (sql.Row2, error) {
+	ri2, ok := i.iter.(sql.RowIter2)
+	if !ok {
+		panic(fmt.Sprintf("%T does not implement sql.RowIter2 interface", i.iter))
+	}
+	row, err := ri2.Next2(ctx)
+	if err != nil {
+		return nil, err
+	}
+	i.numRows++
+	if i.onNext != nil {
+		i.onNext()
+	}
+	return row, nil
+}
+
+func (i *TrackedRowIter) IsRowIter2(ctx *sql.Context) bool {
+	if ri2, ok := i.iter.(sql.RowIter2); ok {
+		return ri2.IsRowIter2(ctx)
+	}
+	return false
+}
+
 func (i *TrackedRowIter) Close(ctx *sql.Context) error {
 	err := i.iter.Close(ctx)
 

sql/rowexec/transaction_iters.go

@@ -15,6 +15,7 @@
 package rowexec
 
 import (
+	"fmt"
 	"io"
 
 	"gopkg.in/src-d/go-errors.v1"
@@ -99,6 +100,21 @@ func (t *TransactionCommittingIter) Next(ctx *sql.Context) (sql.Row, error) {
 	return t.childIter.Next(ctx)
 }
 
+func (t *TransactionCommittingIter) Next2(ctx *sql.Context) (sql.Row2, error) {
+	ri2, ok := t.childIter.(sql.RowIter2)
+	if !ok {
+		panic(fmt.Sprintf("%T does not implement sql.RowIter2 interface", t.childIter))
+	}
+	return ri2.Next2(ctx)
+}
+
+func (t *TransactionCommittingIter) IsRowIter2(ctx *sql.Context) bool {
+	if ri2, ok := t.childIter.(sql.RowIter2); ok {
+		return ri2.IsRowIter2(ctx)
+	}
+	return false
+}
+
 func (t *TransactionCommittingIter) Close(ctx *sql.Context) error {
 	var err error
 	if t.childIter != nil {

sql/rows.go

@@ -92,6 +92,12 @@ type RowIter interface {
 	Closer
 }
 
+type RowIter2 interface {
+	RowIter
+	Next2(ctx *Context) (Row2, error)
+	IsRowIter2(ctx *Context) bool
+}
+
 // RowIterToRows converts a row iterator to a slice of rows.
 func RowIterToRows(ctx *Context, i RowIter) ([]Row, error) {
 	var rows []Row
@@ -112,7 +118,7 @@ func RowIterToRows(ctx *Context, i RowIter) ([]Row, error) {
 	return rows, i.Close(ctx)
 }
 
-func rowFromRow2(sch Schema, r Row2) Row {
+func RowFromRow2(sch Schema, r Row2) Row {
 	row := make(Row, len(sch))
 	for i, col := range sch {
 		switch col.Type.Type() {

sql/table_iter.go

@@ -15,6 +15,7 @@
 package sql
 
 import (
+	"fmt"
 	"io"
 )
 
@@ -24,6 +25,8 @@ type TableRowIter struct {
 	partitions PartitionIter
 	partition  Partition
 	rows       RowIter
+
+	rows2 RowIter2
 }
 
 var _ RowIter = (*TableRowIter)(nil)
@@ -76,6 +79,73 @@ func (i *TableRowIter) Next(ctx *Context) (Row, error) {
 	return row, err
 }
 
+func (i *TableRowIter) Next2(ctx *Context) (Row2, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+	if i.partition == nil {
+		partition, err := i.partitions.Next(ctx)
+		if err != nil {
+			if err == io.EOF {
+				if e := i.partitions.Close(ctx); e != nil {
+					return nil, e
+				}
+			}
+
+			return nil, err
+		}
+
+		i.partition = partition
+	}
+
+	if i.rows2 == nil {
+		rows, err := i.table.PartitionRows(ctx, i.partition)
+		if err != nil {
+			return nil, err
+		}
+		ri2, ok := rows.(RowIter2)
+		if !ok {
+			panic(fmt.Sprintf("%T does not implement RowIter2", rows))
+		}
+		i.rows2 = ri2
+	}
+
+	row, err := i.rows2.Next2(ctx)
+	if err != nil && err == io.EOF {
+		if err = i.rows2.Close(ctx); err != nil {
+			return nil, err
+		}
+		i.partition = nil
+		i.rows2 = nil
+		row, err = i.Next2(ctx)
+	}
+	return row, err
+}
+
+func (i *TableRowIter) IsRowIter2(ctx *Context) bool {
+	if i.partition == nil {
+		partition, err := i.partitions.Next(ctx)
+		if err != nil {
+			return false
+		}
+		i.partition = partition
+	}
+	if i.rows2 == nil {
+		rows, err := i.table.PartitionRows(ctx, i.partition)
+		if err != nil {
+			return false
+		}
+		ri2, ok := rows.(RowIter2)
+		if !ok {
+			return false
+		}
+		i.rows2 = ri2
+	}
+	return i.rows2.IsRowIter2(ctx)
+}
+
 func (i *TableRowIter) Close(ctx *Context) error {
 	if i.rows != nil {
 		if err := i.rows.Close(ctx); err != nil {

sql/type.go

@@ -294,7 +294,6 @@ func IsDecimalType(t Type) bool {
 
 type Type2 interface {
 	Type
-
 	// Compare2 returns an integer comparing two Values.
 	Compare2(Value, Value) (int, error)
 	// Convert2 converts a value of a compatible type.