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

Author: Flak34

CHANGELOG.md

@@ -1,4 +1,5 @@
 * Supported `sql.Null*` from `database/sql` as query params in `toValue` func
+* Added `WithConcurrentResultSets` option for `db.Query().Query()`
 
 ## v3.118.0
 * Added support for nullable `Date32`, `Datetime64`, `Timestamp64`, and `Interval64` types in the `optional` parameter builder

internal/query/client.go

@@ -404,6 +404,7 @@ func clientQuery(ctx context.Context, pool sessionPool, q string, opts ...option
 		if err != nil {
 			return xerrors.WithStackTrace(err)
 		}
+
 		defer func() {
 			_ = streamResult.Close(ctx)
 		}()

internal/query/client_test.go

@@ -922,6 +922,24 @@ func TestClient(t *testing.T) {
 					Status:         Ydb.StatusIds_SUCCESS,
 					ResultSetIndex: 0,
 					ResultSet: &Ydb.ResultSet{
+						Columns: []*Ydb.Column{
+							{
+								Name: "a",
+								Type: &Ydb.Type{
+									Type: &Ydb.Type_TypeId{
+										TypeId: Ydb.Type_UINT64,
+									},
+								},
+							},
+							{
+								Name: "b",
+								Type: &Ydb.Type{
+									Type: &Ydb.Type_TypeId{
+										TypeId: Ydb.Type_UTF8,
+									},
+								},
+							},
+						},
 						Rows: []*Ydb.Value{
 							{
 								Items: []*Ydb.Value{{

internal/query/result.go

@@ -8,12 +8,14 @@ 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_Issue"
 	"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"
@@ -433,34 +435,59 @@ func exactlyOneResultSetFromResult(ctx context.Context, r result.Result) (rs res
 	return MaterializedResultSet(rs.Index(), rs.Columns(), rs.ColumnTypes(), rows), nil
 }
 
-func resultToMaterializedResult(ctx context.Context, r result.Result) (result.Result, error) {
-	var resultSets []result.Set
+func resultToMaterializedResult(ctx context.Context, r *streamResult) (result.Result, error) {
+	type resultSet struct {
+		rows    []query.Row
+		columns []*Ydb.Column
+	}
+	resultSetByIndex := make(map[int64]resultSet)
 
 	for {
-		rs, err := r.NextResultSet(ctx)
+		if ctx.Err() != nil {
+			return nil, xerrors.WithStackTrace(ctx.Err())
+		}
+		if r.closer.Err() != nil {
+			return nil, xerrors.WithStackTrace(r.closer.Err())
+		}
+
+		rs := resultSetByIndex[r.lastPart.GetResultSetIndex()]
+		if len(rs.columns) == 0 {
+			rs.columns = r.lastPart.GetResultSet().GetColumns()
+		}
+
+		rows := make([]query.Row, len(r.lastPart.GetResultSet().GetRows()))
+		for i := range r.lastPart.GetResultSet().GetRows() {
+			rows[i] = NewRow(rs.columns, r.lastPart.GetResultSet().GetRows()[i])
+		}
+		rs.rows = append(rs.rows, rows...)
+
+		resultSetByIndex[r.lastPart.GetResultSetIndex()] = rs
+
+		var err error
+		r.lastPart, err = r.nextPart(ctx)
 		if err != nil {
 			if xerrors.Is(err, io.EOF) {
 				break
 			}
 
 			return nil, xerrors.WithStackTrace(err)
 		}
+		if r.lastPart.GetExecStats() != nil && r.statsCallback != nil {
+			r.statsCallback(stats.FromQueryStats(r.lastPart.GetExecStats()))
+		}
+	}
 
-		var rows []query.Row
-		for {
-			row, err := rs.NextRow(ctx)
-			if err != nil {
-				if xerrors.Is(err, io.EOF) {
-					break
-				}
-
-				return nil, xerrors.WithStackTrace(err)
-			}
+	resultSets := make([]result.Set, len(resultSetByIndex))
+	for rsIndex, rs := range resultSetByIndex {
+		columnNames := make([]string, len(rs.columns))
+		columnTypes := make([]types.Type, len(rs.columns))
 
-			rows = append(rows, row)
+		for i := range rs.columns {
+			columnNames[i] = rs.columns[i].Name
+			columnTypes[i] = types.TypeFromYDB(rs.columns[i].Type)
 		}
 
-		resultSets = append(resultSets, MaterializedResultSet(rs.Index(), rs.Columns(), rs.ColumnTypes(), rows))
+		resultSets[rsIndex] = MaterializedResultSet(int(rsIndex), columnNames, columnTypes, rs.rows)
 	}
 
 	return &materializedResult{

query/execute_options.go

@@ -70,9 +70,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...)