Fb 2048 (#2363)

* Add test coverage for executionplanner.go
*ExecutionPlanner.mapper does not get tested in the case where its
context gets cancelled. In order to make testing this possible,
I've added a context argument to sql_test.MustQueryRow. If it's
nil, MustQueryRow creates a context for itself just like it always
has, but if a context is provided, it uses that.

* Adds test coverage for ExecutionPlanner.mapper in executionplanner.go
The case where the context gets cancelled mid-query is now covered.
The test is timing-dependent - the cancel call has to happen after
the query has been started but before it finishes, and in just the
right part of MustRunQuery, in order to actually produce a context
cancelled error, and not, say, a query cancelled error. May have to
adjust timing if the current delays don't work in CI testing.

* Addressed review notes
-reordered arguments for MustRunQuery
-moved MustRunQuery out of a goroutine, put the cancel in one

sql3/sql_test.go

@@ -46,13 +46,13 @@ func TestSQL_Execute(t *testing.T) {
 
 			// Create a table with all field types.
 			if test.HasTable() {
-				_, _, _, err := sql_test.MustQueryRows(t, svr, test.CreateTable())
+				_, _, _, err := sql_test.MustQueryRows(t, nil, svr, test.CreateTable())
 				assert.NoError(t, err)
 			}
 
 			if test.HasTable() && test.HasData() {
 				// Populate fields with data.
-				_, _, _, err := sql_test.MustQueryRows(t, svr, test.InsertInto(t))
+				_, _, _, err := sql_test.MustQueryRows(t, nil, svr, test.InsertInto(t))
 				assert.NoError(t, err)
 			}
 
@@ -61,7 +61,7 @@ func TestSQL_Execute(t *testing.T) {
 					for _, sql := range sqltest.SQLs {
 						t.Run(fmt.Sprintf("sql-%s", sql), func(t *testing.T) {
 							log.Printf("SQL: %s", sql)
-							rows, headers, plan, err := sql_test.MustQueryRows(t, svr, sql)
+							rows, headers, plan, err := sql_test.MustQueryRows(t, nil, svr, sql)
 
 							// Check expected error instead of results.
 							if sqltest.ExpErr != "" {

sql3/test/helpers.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"testing"
+	"time"
 
 	featurebase "github.com/featurebasedb/featurebase/v3"
 	fbcontext "github.com/featurebasedb/featurebase/v3/context"
@@ -14,14 +15,26 @@ import (
 )
 
 // MustQueryRows returns the row results as a slice of []interface{}, along with the columns, the query plan as a []byte or an error.
-func MustQueryRows(tb testing.TB, svr *featurebase.Server, q string) ([][]interface{}, []*featurebase.WireQueryField, []byte, error) {
+func MustQueryRows(tb testing.TB, c context.Context, svr *featurebase.Server, q string) ([][]interface{}, []*featurebase.WireQueryField, []byte, error) {
 	tb.Helper()
 	requestId, err := uuid.NewV4()
 	if err != nil {
 		return nil, nil, nil, err
 	}
 
-	ctx := fbcontext.WithRequestID(context.Background(), requestId.String())
+	// Originally MustQueryRows just created a context for itself do test with.
+	// However, for some tests, we may want access to the test's context so that,
+	// for example, we can cancel the context mid-test and make sure that gets
+	// handled correctly. Since we need to be able to cancel the context before
+	// the query finishes, if a context is set we also introduce a delay.
+	var ctx context.Context
+	delay := 0 * time.Millisecond
+	if c == nil {
+		ctx = fbcontext.WithRequestID(context.Background(), requestId.String())
+	} else {
+		ctx = fbcontext.WithRequestID(c, requestId.String())
+		delay = 100 * time.Millisecond
+	}
 
 	stmt, err := svr.CompileExecutionPlan(ctx, q)
 	if err != nil {
@@ -43,10 +56,17 @@ func MustQueryRows(tb testing.TB, svr *featurebase.Server, q string) ([][]interf
 	}
 	results := make([][]interface{}, 0)
 
+	// figuring out where to put the sleep took some trial and error.
+	// Too early and the context gets cancelled before the query can
+	// start running, too late and you end up with the query getting
+	// cancelled instead of the context.
+	time.Sleep(delay)
+
 	next, err := rowIter.Next(ctx)
 	if err != nil && err != plannertypes.ErrNoMoreRows {
 		return nil, nil, nil, err
 	}
+
 	for err != plannertypes.ErrNoMoreRows {
 		result := make([]interface{}, len(ocolumns))
 		for i := range result {