feat(tsdb): Add CursorType to enable selection of batch cursors

tsdb/cursor.go

@@ -94,3 +94,31 @@ func CreateCursorIterators(ctx context.Context, shards []*Shard) (CursorIterator
 	}
 	return q, nil
 }
+
+// TODO(sgc): will be removed once batch cursors are gone
+type ctxKey int
+
+const (
+	cursorTypeKey ctxKey = iota
+)
+
+type CursorType int
+
+const (
+	ArrayCursorType CursorType = iota
+	BatchCursorType
+	DefaultCursorType
+)
+
+// NewContextWithCursorType returns a new context with the specified CursorType.
+func NewContextWithCursorType(ctx context.Context, t CursorType) context.Context {
+	return context.WithValue(ctx, cursorTypeKey, t)
+}
+
+// CursorTypeFromContext returns the CursorType associated with ctx or DefaultCursorType if none was set.
+func CursorTypeFromContext(ctx context.Context) CursorType {
+	if v, ok := ctx.Value(cursorTypeKey).(CursorType); ok {
+		return v
+	}
+	return DefaultCursorType
+}

tsdb/engine/tsm1/engine_cursor.go

@@ -2,11 +2,18 @@ package tsm1
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/influxdata/influxdb/tsdb"
 )
 
 func (e *Engine) CreateCursorIterator(ctx context.Context) (tsdb.CursorIterator, error) {
-	return &arrayCursorIterator{e: e}, nil
-	//return &cursorIterator{e: e}, nil
+	switch ct := tsdb.CursorTypeFromContext(ctx); ct {
+	case tsdb.BatchCursorType:
+		return &cursorIterator{e: e}, nil
+	case tsdb.ArrayCursorType, tsdb.DefaultCursorType:
+		return &arrayCursorIterator{e: e}, nil
+	default:
+		panic(fmt.Sprintf("unexpected cursor type %d", ct))
+	}
 }