From c87927a40e26b75af854f27139f374498241eaee Mon Sep 17 00:00:00 2001
From: Andy Arthur <andy@liquidata.co>
Date: Thu, 18 Jun 2020 17:09:14 -0500
Subject: [PATCH] cleanup

---
 bats/query-diff.bats                          |  3 ++
 .../doltcore/diff/querydiff/analyze.go        | 31 ++++++++++++++-----
 .../doltcore/diff/querydiff/query_differ.go   |  6 ++--
 3 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/bats/query-diff.bats b/bats/query-diff.bats
index 2fcd92160a..f0e9977697 100644
--- a/bats/query-diff.bats
+++ b/bats/query-diff.bats
@@ -30,11 +30,14 @@ teardown() {
     dolt sql -q 'insert into test values (0,0,"0"), (1,1,"1")'
     dolt add .
     dolt commit -m rows
+    dolt sql -q 'update test set c1 = 9 where pk = 0'
     dolt sql -q 'delete from test where pk=1'
     dolt sql -q 'insert into test values (2,2,"2")'
     run dolt query_diff 'select * from test'
     [ "$status" -eq 0 ]
     [[ "$output" =~ "|     | pk | c1 | c2 |" ]]
+    [[ "$output" =~ "|  <  | 1  | 1  | 1  |" ]]
+    [[ "$output" =~ "|  >  | 1  | 9  | 1  |" ]]
     [[ "$output" =~ "|  -  | 1  | 1  | 1  |" ]]
     [[ "$output" =~ "|  +  | 2  | 2  | 2  |" ]]
 }
diff --git a/go/libraries/doltcore/diff/querydiff/analyze.go b/go/libraries/doltcore/diff/querydiff/analyze.go
index 351770fceb..7d2fc88e67 100644
--- a/go/libraries/doltcore/diff/querydiff/analyze.go
+++ b/go/libraries/doltcore/diff/querydiff/analyze.go
@@ -51,28 +51,25 @@ func lazyQueryPlan(node sql.Node) (lazyNode sql.Node, projections []sql.Expressi
 	}
 
 	offset := 0
-	newChildren := make([]sql.Node, len(children))
+	lazyChildren := make([]sql.Node, len(children))
 	for i, c := range children {
 		c, pjs, err := lazyQueryPlan(c)
 		if err != nil {
 			return nil, nil, err
 		}
-		if len(c.Schema()) != len(pjs) {
-			return nil, nil, fmt.Errorf("node schema mismatched with node projections")
-		}
-		newChildren[i] = c
+		lazyChildren[i] = c
 
 		pjs = shiftFieldIndexes(pjs, offset)
 		projections = append(projections, pjs...)
 		offset += len(pjs)
 	}
 
-	lazyNode, err = node.WithChildren(newChildren...)
+	node, err = node.WithChildren(lazyChildren...)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	lazyNode, err = plan.TransformExpressions(lazyNode, func(e sql.Expression) (sql.Expression, error) {
+	lazyNode, err = plan.TransformExpressions(node, func(e sql.Expression) (sql.Expression, error) {
 		return makeExpressionLazy(e, projections)
 	})
 	if err != nil {
@@ -84,7 +81,12 @@ func lazyQueryPlan(node sql.Node) (lazyNode sql.Node, projections []sql.Expressi
 		projections = p.Expressions()
 	}
 
-	return lazyNode, projections, err
+	err = validateProjections(lazyNode.Schema(), projections)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return lazyNode, projections, nil
 }
 
 func extractRowOrder(node sql.Node) (composite []plan.SortField) {
@@ -150,6 +152,19 @@ func shiftFieldIndexes(composite []sql.Expression, offset int) []sql.Expression
 	return shifted
 }
 
+func validateProjections(sch sql.Schema, projections []sql.Expression) error {
+	if len(sch) != len(projections) {
+		return fmt.Errorf("lazy node schema does not match lazy projections")
+	}
+	for i, col := range sch {
+		if col.Type != projections[i].Type() {
+			return fmt.Errorf("lazy node schema col type %s does not match lazy projection type %s",
+				col.Type.String(), projections[i].Type().String())
+		}
+	}
+	return nil
+}
+
 func errWithQueryPlan(ctx *sql.Context, eng *sqle.Engine, query string, cause error) error {
 	_, iter, err := eng.Query(ctx, fmt.Sprintf("describe %s", query))
 	if err != nil {
diff --git a/go/libraries/doltcore/diff/querydiff/query_differ.go b/go/libraries/doltcore/diff/querydiff/query_differ.go
index c8edeb6e45..5c6763b559 100644
--- a/go/libraries/doltcore/diff/querydiff/query_differ.go
+++ b/go/libraries/doltcore/diff/querydiff/query_differ.go
@@ -16,7 +16,6 @@ package querydiff
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"io"
 	"math"
@@ -39,8 +38,6 @@ const (
 	Unknown rowOrder = math.MaxInt8
 )
 
-var errSkip = errors.New("errSkip") // u lyk hax?
-
 type rowOrder int8
 
 type QueryDiffer struct {
@@ -108,6 +105,9 @@ func (qd *QueryDiffer) Start() {
 
 func (qd *QueryDiffer) NextDiff() (fromRow sql.Row, toRow sql.Row, err error) {
 	for {
+		if qd.ae.IsSet() {
+			return nil, nil, qd.ae.Get()
+		}
 		if qd.from.isDone() && qd.to.isDone() {
 			return nil, nil, io.EOF
 		}