fix query multiple rowsets (#8738)

ydb/core/viewer/json_pipe_req.h

@@ -150,6 +150,12 @@ protected:
         TString GetError() const {
             return std::get<TString>(Response);
         }
+
+        void Event(const TString& name) {
+            if (Span) {
+                Span.Event(name);
+            }
+        }
     };
 
     NTabletPipe::TClientConfig GetPipeClientConfig();

ydb/core/viewer/viewer_query.h

@@ -21,7 +21,7 @@ class TJsonQuery : public TViewerPipeClient {
     using TBase = TViewerPipeClient;
     TJsonSettings JsonSettings;
     ui32 Timeout = 0;
-    TVector<Ydb::ResultSet> ResultSets;
+    std::vector<std::vector<Ydb::ResultSet>> ResultSets;
     TString Query;
     TString Database;
     TString Action;
@@ -177,11 +177,9 @@ public:
         if (SessionId) {
             auto event = std::make_unique<NKqp::TEvKqp::TEvCloseSessionRequest>();
             event->Record.MutableRequest()->SetSessionId(SessionId);
-            BLOG_TRACE("Closing session " << SessionId);
             Send(NKqp::MakeKqpProxyID(SelfId().NodeId()), event.release());
         }
         TBase::PassAway();
-        BLOG_TRACE("PassAway()");
     }
 
     void ReplyAndPassAway() override {
@@ -227,7 +225,6 @@ public:
                 Span.Attribute("database", Database);
             }
         }
-        BLOG_TRACE("Creating session");
         CreateSessionResponse = MakeRequest<NKqp::TEvKqp::TEvCreateSessionResponse>(NKqp::MakeKqpProxyID(SelfId().NodeId()), event.release());
     }
 
@@ -257,7 +254,6 @@ public:
                     TStringBuilder() << "Failed to create session, error " << ev->Get()->Record.GetYdbStatus()), "InternalError");
         }
         SessionId = CreateSessionResponse->Record.GetResponse().GetSessionId();
-        BLOG_TRACE("Session created " << SessionId);
 
         {
             auto event = std::make_unique<NKqp::TEvKqp::TEvPingSessionRequest>();
@@ -332,7 +328,6 @@ public:
         }
         ActorIdToProto(SelfId(), event->Record.MutableRequestActorId());
         QueryResponse = MakeRequest<NKqp::TEvKqp::TEvQueryResponse>(NKqp::MakeKqpProxyID(SelfId().NodeId()), event.Release());
-        BLOG_TRACE("Query sent");
     }
 
 private:
@@ -432,7 +427,6 @@ private:
     }
 
     void HandleReply(NKqp::TEvKqp::TEvQueryResponse::TPtr& ev) {
-        BLOG_TRACE("Query response received");
         NJson::TJsonValue jsonResponse;
         jsonResponse["version"] = Viewer->GetCapabilityVersion("/viewer/query");
         if (ev->Get()->Record.GetRef().GetYdbStatus() == Ydb::StatusIds::SUCCESS) {
@@ -491,6 +485,7 @@ private:
     }
 
     void HandleReply(NKqp::TEvKqpExecuter::TEvStreamData::TPtr& ev) {
+        QueryResponse.Event("StreamData");
         NKikimrKqp::TEvExecuterStreamData& data(ev->Get()->Record);
 
         if (TotalRows < LimitRows) {
@@ -500,7 +495,10 @@ private:
                 data.MutableResultSet()->set_truncated(true);
             }
             TotalRows += data.GetResultSet().rows_size();
-            ResultSets.emplace_back() = std::move(*data.MutableResultSet());
+            if (ResultSets.size() <= data.GetQueryResultIndex()) {
+                ResultSets.resize(data.GetQueryResultIndex() + 1);
+            }
+            ResultSets[data.GetQueryResultIndex()].emplace_back() = std::move(*data.MutableResultSet());
         }
 
         THolder<NKqp::TEvKqpExecuter::TEvStreamDataAck> ack = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>();
@@ -517,7 +515,6 @@ private:
         if (SessionId) {
             auto event = std::make_unique<NKqp::TEvKqp::TEvCancelQueryRequest>();
             event->Record.MutableRequest()->SetSessionId(SessionId);
-            BLOG_TRACE("Cancelling query in session " << SessionId);
             Send(NKqp::MakeKqpProxyID(SelfId().NodeId()), event.release());
             error << ", query was cancelled";
         }
@@ -549,11 +546,11 @@ private:
                 for (const auto& result : response.GetResults()) {
                     Ydb::ResultSet resultSet;
                     NKqp::ConvertKqpQueryResultToDbResult(result, &resultSet);
-                    ResultSets.emplace_back(std::move(resultSet));
+                    ResultSets.emplace_back().emplace_back(std::move(resultSet));
                 }
 
                 for (const auto& result : response.GetYdbResults()) {
-                    ResultSets.emplace_back(result);
+                    ResultSets.emplace_back().emplace_back(result);
                 }
             }
             catch (const std::exception& ex) {
@@ -568,15 +565,16 @@ private:
             if (Schema == ESchemaType::Classic) {
                 NJson::TJsonValue& jsonResults = jsonResponse["result"];
                 jsonResults.SetType(NJson::JSON_ARRAY);
-                for (auto it = ResultSets.begin(); it != ResultSets.end(); ++it) {
-                    NYdb::TResultSet resultSet(*it);
-                    const auto& columnsMeta = resultSet.GetColumnsMeta();
-                    NYdb::TResultSetParser rsParser(resultSet);
-                    while (rsParser.TryNextRow()) {
-                        NJson::TJsonValue& jsonRow = jsonResults.AppendValue({});
-                        for (size_t columnNum = 0; columnNum < columnsMeta.size(); ++columnNum) {
-                            const NYdb::TColumn& columnMeta = columnsMeta[columnNum];
-                            jsonRow[columnMeta.Name] = ColumnValueToJsonValue(rsParser.ColumnParser(columnNum));
+                for (const auto& resultSets : ResultSets) {
+                    for (NYdb::TResultSet resultSet : resultSets) {
+                        const auto& columnsMeta = resultSet.GetColumnsMeta();
+                        NYdb::TResultSetParser rsParser(resultSet);
+                        while (rsParser.TryNextRow()) {
+                            NJson::TJsonValue& jsonRow = jsonResults.AppendValue({});
+                            for (size_t columnNum = 0; columnNum < columnsMeta.size(); ++columnNum) {
+                                const NYdb::TColumn& columnMeta = columnsMeta[columnNum];
+                                jsonRow[columnMeta.Name] = ColumnValueToJsonValue(rsParser.ColumnParser(columnNum));
+                            }
                         }
                     }
                 }
@@ -585,7 +583,7 @@ private:
             if (Schema == ESchemaType::Modern) {
                 {
                     NJson::TJsonValue& jsonColumns = jsonResponse["columns"];
-                    NYdb::TResultSet resultSet(ResultSets.front());
+                    NYdb::TResultSet resultSet(ResultSets.front().front());
                     const auto& columnsMeta = resultSet.GetColumnsMeta();
                     jsonColumns.SetType(NJson::JSON_ARRAY);
                     for (size_t columnNum = 0; columnNum < columnsMeta.size(); ++columnNum) {
@@ -598,16 +596,17 @@ private:
 
                 NJson::TJsonValue& jsonResults = jsonResponse["result"];
                 jsonResults.SetType(NJson::JSON_ARRAY);
-                for (auto it = ResultSets.begin(); it != ResultSets.end(); ++it) {
-                    NYdb::TResultSet resultSet(*it);
-                    const auto& columnsMeta = resultSet.GetColumnsMeta();
-                    NYdb::TResultSetParser rsParser(resultSet);
-                    while (rsParser.TryNextRow()) {
-                        NJson::TJsonValue& jsonRow = jsonResults.AppendValue({});
-                        jsonRow.SetType(NJson::JSON_ARRAY);
-                        for (size_t columnNum = 0; columnNum < columnsMeta.size(); ++columnNum) {
-                            NJson::TJsonValue& jsonColumn = jsonRow.AppendValue({});
-                            jsonColumn = ColumnValueToJsonValue(rsParser.ColumnParser(columnNum));
+                for (const auto& resultSets : ResultSets) {
+                    for (NYdb::TResultSet resultSet : resultSets) {
+                        const auto& columnsMeta = resultSet.GetColumnsMeta();
+                        NYdb::TResultSetParser rsParser(resultSet);
+                        while (rsParser.TryNextRow()) {
+                            NJson::TJsonValue& jsonRow = jsonResults.AppendValue({});
+                            jsonRow.SetType(NJson::JSON_ARRAY);
+                            for (size_t columnNum = 0; columnNum < columnsMeta.size(); ++columnNum) {
+                                NJson::TJsonValue& jsonColumn = jsonRow.AppendValue({});
+                                jsonColumn = ColumnValueToJsonValue(rsParser.ColumnParser(columnNum));
+                            }
                         }
                     }
                 }
@@ -616,47 +615,53 @@ private:
             if (Schema == ESchemaType::Multi) {
                 NJson::TJsonValue& jsonResults = jsonResponse["result"];
                 jsonResults.SetType(NJson::JSON_ARRAY);
-                for (auto it = ResultSets.begin(); it != ResultSets.end(); ++it) {
-                    NYdb::TResultSet resultSet(*it);
-                    const auto& columnsMeta = resultSet.GetColumnsMeta();
+                for (const auto& resultSets : ResultSets) {
                     NJson::TJsonValue& jsonResult = jsonResults.AppendValue({});
+                    bool hasColumns = false;
+                    for (NYdb::TResultSet resultSet : resultSets) {
+                        if (!hasColumns) {
+                            NJson::TJsonValue& jsonColumns = jsonResult["columns"];
+                            jsonColumns.SetType(NJson::JSON_ARRAY);
+                            const auto& columnsMeta = resultSet.GetColumnsMeta();
+                            for (size_t columnNum = 0; columnNum < columnsMeta.size(); ++columnNum) {
+                                NJson::TJsonValue& jsonColumn = jsonColumns.AppendValue({});
+                                const NYdb::TColumn& columnMeta = columnsMeta[columnNum];
+                                jsonColumn["name"] = columnMeta.Name;
+                                jsonColumn["type"] = columnMeta.Type.ToString();
+                            }
+                            hasColumns = true;
+                        }
 
-                    NJson::TJsonValue& jsonColumns = jsonResult["columns"];
-                    jsonColumns.SetType(NJson::JSON_ARRAY);
-                    for (size_t columnNum = 0; columnNum < columnsMeta.size(); ++columnNum) {
-                        NJson::TJsonValue& jsonColumn = jsonColumns.AppendValue({});
-                        const NYdb::TColumn& columnMeta = columnsMeta[columnNum];
-                        jsonColumn["name"] = columnMeta.Name;
-                        jsonColumn["type"] = columnMeta.Type.ToString();
-                    }
+                        NJson::TJsonValue& jsonRows = jsonResult["rows"];
+                        NYdb::TResultSetParser rsParser(resultSet);
+                        while (rsParser.TryNextRow()) {
+                            NJson::TJsonValue& jsonRow = jsonRows.AppendValue({});
+                            jsonRow.SetType(NJson::JSON_ARRAY);
+                            for (size_t columnNum = 0; columnNum < rsParser.ColumnsCount(); ++columnNum) {
+                                NJson::TJsonValue& jsonColumn = jsonRow.AppendValue({});
+                                jsonColumn = ColumnValueToJsonValue(rsParser.ColumnParser(columnNum));
+                            }
+                        }
 
-                    NJson::TJsonValue& jsonRows = jsonResult["rows"];
-                    NYdb::TResultSetParser rsParser(resultSet);
-                    while (rsParser.TryNextRow()) {
-                        NJson::TJsonValue& jsonRow = jsonRows.AppendValue({});
-                        jsonRow.SetType(NJson::JSON_ARRAY);
-                        for (size_t columnNum = 0; columnNum < columnsMeta.size(); ++columnNum) {
-                            NJson::TJsonValue& jsonColumn = jsonRow.AppendValue({});
-                            jsonColumn = ColumnValueToJsonValue(rsParser.ColumnParser(columnNum));
+                        if (resultSet.Truncated()) {
+                            jsonResult["truncated"] = true;
                         }
                     }
-                    if (resultSet.Truncated()) {
-                        jsonResult["truncated"] = true;
-                    }
                 }
             }
 
             if (Schema == ESchemaType::Ydb) {
                 NJson::TJsonValue& jsonResults = jsonResponse["result"];
                 jsonResults.SetType(NJson::JSON_ARRAY);
-                for (auto it = ResultSets.begin(); it != ResultSets.end(); ++it) {
-                    NYdb::TResultSet resultSet(*it);
-                    const auto& columnsMeta = resultSet.GetColumnsMeta();
-                    NYdb::TResultSetParser rsParser(resultSet);
-                    while (rsParser.TryNextRow()) {
-                        NJson::TJsonValue& jsonRow = jsonResults.AppendValue({});
-                        TString row = NYdb::FormatResultRowJson(rsParser, columnsMeta, IsBase64Encode ? NYdb::EBinaryStringEncoding::Base64 : NYdb::EBinaryStringEncoding::Unicode);
-                        NJson::ReadJsonTree(row, &jsonRow);
+                for (const auto& resultSets : ResultSets) {
+                    for (NYdb::TResultSet resultSet : resultSets) {
+                        const auto& columnsMeta = resultSet.GetColumnsMeta();
+                        NYdb::TResultSetParser rsParser(resultSet);
+                        while (rsParser.TryNextRow()) {
+                            NJson::TJsonValue& jsonRow = jsonResults.AppendValue({});
+                            TString row = NYdb::FormatResultRowJson(rsParser, columnsMeta, IsBase64Encode ? NYdb::EBinaryStringEncoding::Base64 : NYdb::EBinaryStringEncoding::Unicode);
+                            NJson::ReadJsonTree(row, &jsonRow);
+                        }
                     }
                 }
             }