Add concurrent result sets in db.Query.Query(...)

Author: Flak34

internal/query/client.go

@@ -397,15 +397,23 @@ func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...option
 ) {
 	settings := options.ExecuteSettings(opts...)
 	err = do(ctx, pool, func(ctx context.Context, s *Session) (err error) {
-		streamResult, err := s.execute(ctx, q, options.ExecuteSettings(opts...), withStreamResultTrace(s.trace))
+		request, callOptions, err := executeQueryRequest(s.ID(), q, settings)
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
+		request.ConcurrentResultSets = settings.ConcurrentResultSets()
+
+		executeCtx, executeCancel := xcontext.WithCancel(xcontext.ValueOnly(ctx))
 		defer func() {
-			_ = streamResult.Close(ctx)
+			executeCancel()
 		}()
 
-		r, err = resultToMaterializedResult(ctx, streamResult)
+		stream, err := s.client.ExecuteQuery(executeCtx, request, callOptions...)
+		if err != nil {
+			return xerrors.WithStackTrace(err)
+		}
+
+		r, err = streamToMaterializedResult(ctx, stream)
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}

internal/query/execute_query.go

@@ -93,7 +93,7 @@ func executeQueryRequest(sessionID, q string, cfg executeSettings) (
 		},
 		Parameters:             params,
 		StatsMode:              Ydb_Query.StatsMode(cfg.StatsMode()),
-		ConcurrentResultSets:   cfg.ConcurrentResultSets(),
+		ConcurrentResultSets:   false,
 		PoolId:                 cfg.ResourcePool(),
 		ResponsePartLimitBytes: cfg.ResponsePartLimitSizeBytes(),
 	}

internal/query/result.go

@@ -8,11 +8,12 @@ import (
 	"time"
 
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
-
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/result"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stats"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/types"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xiter"
 	"github.com/ydb-platform/ydb-go-sdk/v3/query"
@@ -459,3 +460,60 @@ func resultToMaterializedResult(ctx context.Context, r result.Result) (result.Re
 		resultSets: resultSets,
 	}, nil
 }
+
+func streamToMaterializedResult(ctx context.Context, stream Ydb_Query_V1.QueryService_ExecuteQueryClient) (result.Result, error) {
+	type resultSet struct {
+		rows        []query.Row
+		rawColumns  []*Ydb.Column
+		columnNames []string
+		columnTypes []types.Type
+	}
+	resultSetByIndex := make(map[int64]resultSet)
+
+	for {
+		if ctx.Err() != nil {
+			return nil, xerrors.WithStackTrace(ctx.Err())
+		}
+
+		part, err := stream.Recv()
+		if err != nil {
+			if xerrors.Is(err, io.EOF) {
+				break
+			}
+			return nil, xerrors.WithStackTrace(err)
+		}
+
+		if part.GetResultSetIndex() < 0 {
+			break
+		}
+
+		rs := resultSetByIndex[part.GetResultSetIndex()]
+		if len(rs.columnNames) == 0 {
+			rs.rawColumns = part.GetResultSet().GetColumns()
+			rs.columnNames = make([]string, 0, len(rs.rawColumns))
+			rs.columnTypes = make([]types.Type, 0, len(rs.rawColumns))
+
+			for _, column := range rs.rawColumns {
+				rs.columnNames = append(rs.columnNames, column.GetName())
+				rs.columnTypes = append(rs.columnTypes, types.TypeFromYDB(column.GetType()))
+			}
+		}
+
+		rows := make([]query.Row, 0, len(part.GetResultSet().GetRows()))
+		for _, row := range part.GetResultSet().GetRows() {
+			rows = append(rows, NewRow(rs.rawColumns, row))
+		}
+		rs.rows = append(rs.rows, rows...)
+
+		resultSetByIndex[part.GetResultSetIndex()] = rs
+	}
+
+	resultSets := make([]result.Set, len(resultSetByIndex))
+	for rsIndex, rs := range resultSetByIndex {
+		resultSets[rsIndex] = MaterializedResultSet(int(rsIndex), rs.columnNames, rs.columnTypes, rs.rows)
+	}
+
+	return &materializedResult{
+		resultSets: resultSets,
+	}, nil
+}

query/execute_options.go

@@ -63,9 +63,9 @@ func WithResponsePartLimitSizeBytes(size int64) ExecuteOption {
 	return options.WithResponsePartLimitSizeBytes(size)
 }
 
-//func WithConcurrentResultSets(isEnabled bool) ExecuteOption {
-//	return options.WithConcurrentResultSets(isEnabled)
-//}
+func WithConcurrentResultSets(isEnabled bool) ExecuteOption {
+	return options.WithConcurrentResultSets(isEnabled)
+}
 
 func WithCallOptions(opts ...grpc.CallOption) ExecuteOption {
 	return options.WithCallOptions(opts...)