change json/query format

1 files changed, 87 insertions, 52 deletions

diff --git a/ydb/core/viewer/json_query.h b/ydb/core/viewer/json_query.h
index bb07dcc521b..ac32ad0fc5c 100644
--- a/ydb/core/viewer/json_query.h
+++ b/ydb/core/viewer/json_query.h
@@ -8,10 +8,12 @@
 #include <library/cpp/json/json_reader.h>
 #include <library/cpp/protobuf/json/proto2json.h>
 #include <ydb/core/node_whiteboard/node_whiteboard.h>
+#include <ydb/core/grpc_services/rpc_kqp_base.h>
 #include <ydb/core/kqp/kqp.h>
 #include <ydb/core/kqp/executer/kqp_executer.h>
 #include <ydb/core/viewer/json/json.h>
-#include <ydb/public/lib/deprecated/kicli/kicli.h>
+//#include <ydb/public/lib/deprecated/kicli/kicli.h>
+#include <ydb/public/lib/json_value/ydb_json_value.h>
 #include <ydb/public/sdk/cpp/client/ydb_result/result.h>
 
 namespace NKikimr {
@@ -31,6 +33,7 @@ class TJsonQuery : public TActorBootstrapped<TJsonQuery> {
     TVector<Ydb::ResultSet> ResultSets;
     TString Action;
     TString Stats;
+    TString Schema = "classic";
 
 public:
     static constexpr NKikimrServices::TActivity::EType ActorActivityType() {
@@ -70,6 +73,10 @@ public:
         TString database = params.Get("database");
         Stats = params.Get("stats");
         Action = params.Get("action");
+        Schema = params.Get("schema");
+        if (Schema.empty()) {
+            Schema = "classic";
+        }
         if (query.empty() && Event->Get()->Request.GetMethod() == HTTP_METHOD_POST) {
             TStringBuf content = Event->Get()->Request.GetPostContent();
             const THttpHeaders& headers = Event->Get()->Request.GetHeaders();
@@ -205,74 +212,96 @@ private:
 
     void HandleReply(NKqp::TEvKqp::TEvQueryResponse::TPtr& ev) {
         TStringBuilder out;
+        NJson::TJsonValue jsonResponse;
         NKikimrKqp::TEvQueryResponse& record = ev->Get()->Record.GetRef();
         if (record.GetYdbStatus() == Ydb::StatusIds::SUCCESS) {
             const auto& response = record.GetResponse();
-            out << Viewer->GetHTTPOKJSON(Event->Get());
-            if (!Stats.empty()) {
-                out << "{\"result\":";
+
+            if (response.ResultsSize() > 0) {
+                for (const auto& result : response.GetResults()) {
+                    Ydb::ResultSet resultSet;
+                    NGRpcService::ConvertKqpQueryResultToDbResult(result, &resultSet);
+                    ResultSets.emplace_back(std::move(resultSet));
+                }
             }
+
+            out << Viewer->GetHTTPOKJSON(Event->Get());
             if (ResultSets.size() > 0) {
-                out << "[";
-                bool comma = false;
-                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()) {
-                        if (comma) {
-                            out << ",";
+                if (Schema == "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));
+                            }
                         }
-                        out << "{";
+                    }
+                }
+
+                if (Schema == "modern") {
+                    {
+                        NJson::TJsonValue& jsonColumns = jsonResponse["columns"];
+                        NYdb::TResultSet resultSet(ResultSets.front());
+                        const auto& columnsMeta = resultSet.GetColumnsMeta();
+                        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];
-                            out << "\"" << TProtoToJson::EscapeJsonString(columnMeta.Name) << "\":";
-                            out << NJson::WriteJson(ColumnValueToJsonValue(rsParser.ColumnParser(columnNum)), false);
-                            if (columnNum + 1 < columnsMeta.size()) {
-                                out << ",";
+                            jsonColumn["name"] = columnMeta.Name;
+                            jsonColumn["type"] = columnMeta.Type.ToString();
+                        }
+                    }
+
+                    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));
                             }
                         }
-                        out << "}";
-                        comma = true;
                     }
                 }
-                out << ']';
-            } else if (response.ResultsSize() > 0) {
-                const auto &result = response.GetResults();
-                if (result.empty()) {
-                    out << "[]";
-                } else {
-                    const auto &first = *result.begin();
-
-                    if (first.HasType() && first.HasValue()) {
-                        auto value = NClient::TValue::Create(first.GetValue(), first.GetType());
-                        auto data = value["Data"];
-                        if (data.HaveValue()) {
-                            out << data.GetValueText<NClient::TFormatJSON>({JsonSettings.UI64AsString});
-                        } else {
-                            out << "[]";
+
+                if (Schema == "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, NYdb::EBinaryStringEncoding::Base64);
+                            NJson::ReadJsonTree(row, &jsonRow);
                         }
-                    } else {
-                        out << "[]";
                     }
                 }
-            } else if (response.HasQueryPlan()) {
-                if (Action == "explain-ast") {
-                    out << "{\"ast\":\"" << TProtoToJson::EscapeJsonString(response.GetQueryAst()) << "\"}";
-                } else {
-                    out << response.GetQueryPlan();
-                }
             }
-            if (!Stats.empty()) {
-                out << ",\"stats\":";
-                TStringStream json;
-                TProtoToJson::ProtoToJson(json, response.GetQueryStats(), JsonSettings);
-                out << json.Str() << "}";
+            if (response.HasQueryAst()) {
+                jsonResponse["ast"] = response.GetQueryAst();
+            }
+            if (response.HasQueryPlan()) {
+                NJson::ReadJsonTree(response.GetQueryPlan(), &(jsonResponse["plan"]));
+            }
+            if (response.HasQueryStats()) {
+                NProtobufJson::Proto2Json(response.GetQueryStats(), jsonResponse["stats"]);
             }
         } else {
             out << "HTTP/1.1 400 Bad Request\r\nContent-Type: application/json\r\nConnection: Close\r\n\r\n";
-            NJson::TJsonValue response;
-            NJson::TJsonValue& jsonIssues = response["issues"];
+            NJson::TJsonValue& jsonIssues = jsonResponse["issues"];
 
             // find first deepest error
             google::protobuf::RepeatedPtrField<Ydb::Issue::IssueMessage>* protoIssues = record.MutableResponse()->MutableQueryIssues();
@@ -284,15 +313,20 @@ private:
             }
             if (protoIssues->size() > 0) {
                 const Ydb::Issue::IssueMessage& issue = (*protoIssues)[0];
-                NProtobufJson::Proto2Json(issue, response["error"]);
+                NProtobufJson::Proto2Json(issue, jsonResponse["error"]);
             }
             for (const auto& queryIssue : record.GetResponse().GetQueryIssues()) {
                 NJson::TJsonValue& issue = jsonIssues.AppendValue({});
                 NProtobufJson::Proto2Json(queryIssue, issue);
             }
-            out << NJson::WriteJson(response, false);
         }
 
+        if (Schema == "classic" && Stats.empty()) {
+            jsonResponse = std::move(jsonResponse["result"]);
+        }
+
+        out << NJson::WriteJson(jsonResponse, false);
+
         ReplyAndPassAway(out);
     }
 
@@ -339,6 +373,7 @@ struct TJsonRequestParameters<TJsonQuery> {
         return R"___([{"name":"ui64","in":"query","description":"return ui64 as number","required":false,"type":"boolean"},
                       {"name":"query","in":"query","description":"query text","required":true,"type":"string"},
                       {"name":"database","in":"query","description":"database name","required":false,"type":"string"},
+                      {"name":"schema","in":"query","description":"result format schema (classic, modern, ydb)","required":false,"type":"string"},
                       {"name":"stats","in":"query","description":"return stats (profile)","required":false,"type":"string"},
                       {"name":"action","in":"query","description":"execute method (execute-scan, execute-script, explain, explain-ast)","required":false,"type":"string"},
                       {"name":"timeout","in":"query","description":"timeout in ms","required":false,"type":"integer"}])___";