peek through limit iters as well

server/handler.go

@@ -40,6 +40,7 @@ import (
 	"github.com/dolthub/go-mysql-server/internal/sockstate"
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/analyzer"
+	"github.com/dolthub/go-mysql-server/sql/iters"
 	"github.com/dolthub/go-mysql-server/sql/plan"
 	"github.com/dolthub/go-mysql-server/sql/rowexec"
 	"github.com/dolthub/go-mysql-server/sql/types"
@@ -515,21 +516,33 @@ func resultForEmptyIter(ctx *sql.Context, iter sql.RowIter, resultFields []*quer
 
 // GetDeferredProjections looks for a top-level deferred projection, retrieves its projections, and removes it from the
 // iterator tree.
-func GetDeferredProjections(iter sql.RowIter) []sql.Expression {
+func GetDeferredProjections(iter sql.RowIter) (sql.RowIter, []sql.Expression) {
 	switch i := iter.(type) {
 	case *rowexec.ExprCloserIter:
-		return GetDeferredProjections(i.GetIter())
+		_, projs := GetDeferredProjections(i.GetIter())
+		return i, projs
 	case *plan.TrackedRowIter:
-		if commitIter, isCommitIter := i.GetIter().(*rowexec.TransactionCommittingIter); isCommitIter {
-			if projIter, isProjIter := commitIter.GetIter().(*rowexec.ProjectIter); isProjIter {
-				if projIter.CanDefer() {
-					commitIter.WithChildIter(projIter.GetChildIter())
-					return projIter.GetProjections()
-				}
-			}
+		_, projs := GetDeferredProjections(i.GetIter())
+		return i, projs
+	case *rowexec.TransactionCommittingIter:
+		newChild, projs := GetDeferredProjections(i.GetIter())
+		if projs != nil {
+			 i.WithChildIter(newChild)
+		}
+		return i, projs
+	case *iters.LimitIter:
+		newChild, projs := GetDeferredProjections(i.ChildIter)
+		if projs != nil {
+			i.ChildIter = newChild
 		}
+		return i, projs
+	case *rowexec.ProjectIter:
+		if i.CanDefer() {
+			return i.GetChildIter(), i.GetProjections()
+		}
+		return i, nil
 	}
-	return nil
+	return iter, nil
 }
 
 // resultForMax1RowIter ensures that an empty iterator returns at most one row
@@ -582,7 +595,7 @@ func (h *Handler) resultForDefaultIter(
 	wg.Add(2)
 
 	// Read rows off the row iterator and send them to the row channel.
-	projs := GetDeferredProjections(iter)
+	iter, projs := GetDeferredProjections(iter)
 	var rowChan = make(chan sql.Row, 512)
 	eg.Go(func() error {
 		defer pan2err()