diff --git a/CHANGELOG.md b/CHANGELOG.md index bb0357e37..b1f66a7f4 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/internal/query/client.go b/internal/query/client.go index 62a93c1e5..f89ffde41 100644 --- a/internal/query/client.go +++ b/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) }() diff --git a/internal/query/client_test.go b/internal/query/client_test.go index 2421f6259..a87911d79 100644 --- a/internal/query/client_test.go +++ b/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{{ diff --git a/internal/query/execute_query.go b/internal/query/execute_query.go index 6cdcb622c..243d8ad1a 100644 --- a/internal/query/execute_query.go +++ b/internal/query/execute_query.go @@ -32,6 +32,7 @@ type executeSettings interface { ResourcePool() string ResponsePartLimitSizeBytes() int64 Label() string + ConcurrentResultSets() bool } type executeScriptConfig interface { @@ -92,7 +93,7 @@ func executeQueryRequest(sessionID, q string, cfg executeSettings) ( }, Parameters: params, StatsMode: Ydb_Query.StatsMode(cfg.StatsMode()), - ConcurrentResultSets: false, + ConcurrentResultSets: cfg.ConcurrentResultSets(), PoolId: cfg.ResourcePool(), ResponsePartLimitBytes: cfg.ResponsePartLimitSizeBytes(), } diff --git a/internal/query/options/execute.go b/internal/query/options/execute.go index 6853c02d9..4e272c601 100644 --- a/internal/query/options/execute.go +++ b/internal/query/options/execute.go @@ -46,6 +46,7 @@ type ( issueCallback func(issues []*Ydb_Issue.IssueMessage) responsePartLimitBytes int64 label string + concurrentResultSets bool } // Execute is an interface for execute method options @@ -72,9 +73,12 @@ type ( } execModeOption = ExecMode responsePartLimitBytes int64 - issuesOption struct { + + issuesOption struct { callback func([]*Ydb_Issue.IssueMessage) } + + concurrentResultSets bool ) func (poolID resourcePool) applyExecuteOption(s *executeSettings) { @@ -132,6 +136,10 @@ func (opts issuesOption) applyExecuteOption(s *executeSettings) { s.issueCallback = opts.callback } +func (opt concurrentResultSets) applyExecuteOption(s *executeSettings) { + s.concurrentResultSets = bool(opt) +} + const ( ExecModeParse = ExecMode(Ydb_Query.ExecMode_EXEC_MODE_PARSE) ExecModeValidate = ExecMode(Ydb_Query.ExecMode_EXEC_MODE_VALIDATE) @@ -205,6 +213,10 @@ func (s *executeSettings) Label() string { return s.label } +func (s *executeSettings) ConcurrentResultSets() bool { + return s.concurrentResultSets +} + func WithParameters(params params.Parameters) parametersOption { return parametersOption{ params: params, @@ -237,6 +249,10 @@ func WithResponsePartLimitSizeBytes(size int64) responsePartLimitBytes { return responsePartLimitBytes(size) } +func WithConcurrentResultSets(isEnabled bool) concurrentResultSets { + return concurrentResultSets(isEnabled) +} + func (size responsePartLimitBytes) applyExecuteOption(s *executeSettings) { s.responsePartLimitBytes = int64(size) } diff --git a/internal/query/result.go b/internal/query/result.go index b26a97a29..01b9a387b 100644 --- a/internal/query/result.go +++ b/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,11 +435,36 @@ 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 @@ -445,22 +472,22 @@ func resultToMaterializedResult(ctx context.Context, r result.Result) (result.Re 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{ diff --git a/query/execute_options.go b/query/execute_options.go index c0f6cfb0a..3e188dcd7 100644 --- a/query/execute_options.go +++ b/query/execute_options.go @@ -70,6 +70,10 @@ func WithResponsePartLimitSizeBytes(size int64) ExecuteOption { return options.WithResponsePartLimitSizeBytes(size) } +func WithConcurrentResultSets(isEnabled bool) ExecuteOption { + return options.WithConcurrentResultSets(isEnabled) +} + func WithCallOptions(opts ...grpc.CallOption) ExecuteOption { return options.WithCallOptions(opts...) }