[ydb-cli] saving plan & ast even on errors (#7874)

Author: gridnevvvit

ydb/public/lib/ydb_cli/commands/benchmark_utils.cpp

@@ -162,13 +162,6 @@ class IQueryResultScanner {
     bool Scan(TIterator& it) {
         for (;;) {
             auto streamPart = it.ReadNext().GetValueSync();
-            if (!streamPart.IsSuccess()) {
-                if (!streamPart.EOS()) {
-                    OnError(streamPart.GetIssues().ToString());
-                    return false;
-                }
-                break;
-            }
 
             if constexpr (std::is_same_v<TIterator, NTable::TScanQueryPartIterator>) {
                 if (streamPart.HasQueryStats()) {
@@ -185,6 +178,14 @@ class IQueryResultScanner {
                 }
             }
 
+            if (!streamPart.IsSuccess()) {
+                if (!streamPart.EOS()) {
+                    OnError(streamPart.GetIssues().ToString());
+                    return false;
+                }
+                break;
+            }
+
             if (streamPart.HasResultSet()) {
                 auto result = streamPart.ExtractResultSet();
                 auto columns = result.GetColumnsMeta();
@@ -302,7 +303,8 @@ TQueryBenchmarkResult Execute(const TString& query, NTable::TTableClient& client
     composite.AddScanner(scannerYson);
     composite.AddScanner(scannerCSV);
     if (!composite.Scan(it)) {
-        return TQueryBenchmarkResult::Error(composite.GetErrorInfo());
+        return TQueryBenchmarkResult::Error(
+            composite.GetErrorInfo(), composite.GetQueryPlan(), composite.GetPlanAst());
     } else {
         return TQueryBenchmarkResult::Result(
             scannerYson->GetResult(),
@@ -329,7 +331,8 @@ TQueryBenchmarkResult Execute(const TString& query, NQuery::TQueryClient& client
     composite.AddScanner(scannerYson);
     composite.AddScanner(scannerCSV);
     if (!composite.Scan(it)) {
-        return TQueryBenchmarkResult::Error(composite.GetErrorInfo());
+        return TQueryBenchmarkResult::Error(
+            composite.GetErrorInfo(), composite.GetQueryPlan(), composite.GetPlanAst());
     } else {
         return TQueryBenchmarkResult::Result(
             scannerYson->GetResult(),

ydb/public/lib/ydb_cli/commands/benchmark_utils.h

@@ -74,9 +74,11 @@ class TQueryBenchmarkResult {
         result.PlanAst = planAst;
         return result;
     }
-    static TQueryBenchmarkResult Error(const TString& error) {
+    static TQueryBenchmarkResult Error(const TString& error, const TString& queryPlan, const TString& planAst) {
         TQueryBenchmarkResult result;
         result.ErrorInfo = error;
+        result.QueryPlan = queryPlan;
+        result.PlanAst = planAst;
         return result;
     }
     operator bool() const {

ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp

@@ -10,7 +10,7 @@ namespace NYdb::NConsoleClient {
     TWorkloadCommandBenchmark::TWorkloadCommandBenchmark(NYdbWorkload::TWorkloadParams& params, const NYdbWorkload::IWorkloadQueryGenerator::TWorkloadType& workload)
         : TWorkloadCommandBase(workload.CommandName, params, NYdbWorkload::TWorkloadParams::ECommandType::Run, workload.Description, workload.Type)
     {
-        
+
     }
 
 
@@ -333,11 +333,11 @@ bool TWorkloadCommandBenchmark::RunBench(TClient& client, NYdbWorkload::IWorkloa
         bool planSaved = false;
         for (ui32 i = 0; i < IterationsCount; ++i) {
             auto t1 = TInstant::Now();
-            TQueryBenchmarkResult res = TQueryBenchmarkResult::Error("undefined");
+            TQueryBenchmarkResult res = TQueryBenchmarkResult::Error("undefined", "undefined", "undefined");
             try {
                 res = Execute(query, client);
             } catch (...) {
-                res = TQueryBenchmarkResult::Error(CurrentExceptionMessage());
+                res = TQueryBenchmarkResult::Error(CurrentExceptionMessage(), "", "");
             }
             auto duration = TInstant::Now() - t1;