Support progress stats for ydb_cli (hidden option)

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

@@ -169,7 +169,7 @@ class TQueryResultScanner {
     template <typename TIterator>
     TStatus Scan(TIterator& it, std::optional<TString> planFileName = std::nullopt) {
 
-        TProgressIndication progressIndication(true);
+        TProgressIndication progressIndication;
         TMaybe<NQuery::TExecStats> execStats;
 
         TString currentPlanFileNameStats;
@@ -222,11 +222,10 @@ class TQueryResultScanner {
                     const auto& protoStats = TProtoAccessor::GetProto(execStats.GetRef());
                     for (const auto& queryPhase : protoStats.query_phases()) {
                         for (const auto& tableAccessStats : queryPhase.table_access()) {
-                            progressIndication.UpdateProgress({tableAccessStats.reads().rows(), tableAccessStats.reads().bytes(),
-                                tableAccessStats.updates().rows(), tableAccessStats.updates().bytes(),
-                                tableAccessStats.deletes().rows(), tableAccessStats.deletes().bytes()});
+                            progressIndication.UpdateProgress({tableAccessStats.reads().rows(), tableAccessStats.reads().bytes()});
                         }
                     }
+                    progressIndication.SetDurationUs(protoStats.total_duration_us());
 
                     progressIndication.Render();
                 }

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

@@ -6,8 +6,10 @@
 #include <ydb/public/lib/json_value/ydb_json_value.h>
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/library/operation_id/operation_id.h>
 #include <ydb/public/lib/ydb_cli/common/interactive.h>
+#include <ydb/public/lib/ydb_cli/common/plan2svg.h>
 #include <ydb/public/lib/ydb_cli/common/pretty_table.h>
 #include <ydb/public/lib/ydb_cli/common/print_operation.h>
+#include <ydb/public/lib/ydb_cli/common/progress_indication.h>
 #include <ydb/public/lib/ydb_cli/common/query_stats.h>
 #include <ydb/public/lib/ydb_cli/common/waiting_bar.h>
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/proto/accessor.h>
@@ -43,6 +45,18 @@ void TCommandSql::Config(TConfig& config) {
         .StoreTrue(&ExplainAnalyzeMode);
     config.Opts->AddLongOption("stats", "Execution statistics collection mode [none, basic, full, profile]")
         .RequiredArgument("[String]").StoreResult(&CollectStatsMode);
+
+    NColorizer::TColors colors = NColorizer::AutoColors(Cout);
+    TStringStream description;
+    description << "Print progress of query execution. Requires non-none statistics collection mode. Available options: ";
+    description << "\n  " << colors.BoldColor() << "tty" << colors.OldColor()
+            << "\n    " << "Print progress to the terminal";
+    description << "\n  " << colors.BoldColor() << "none" << colors.OldColor()
+            << "\n    " << "Disables progress printing";
+    description << "\nDefault: " << colors.CyanColor() << "\"none\"" << colors.OldColor() << ".";
+
+    config.Opts->AddLongOption("progress", description.Str())
+        .RequiredArgument("[String]").Hidden().DefaultValue("none").StoreResult(&Progress);
     config.Opts->AddLongOption("diagnostics-file", "Path to file where the diagnostics will be saved.")
         .RequiredArgument("[String]").StoreResult(&DiagnosticsFile);
     config.Opts->AddLongOption("syntax", "Query syntax [yql, pg]")
@@ -122,6 +136,9 @@ void TCommandSql::Parse(TConfig& config) {
             << "nor path to file with script text (\"--file\", \"-f\") were provided." << Endl;
         config.PrintHelpAndExit();
     }
+    if (Progress && Progress != "tty" && Progress != "none") {
+        throw TMisuseException() << "Unknow progress option \"" << Progress << "\".";
+    }
     // Should be called after setting ReadingSomethingFromStdin
     ParseParameters(config);
 }
@@ -144,7 +161,18 @@ int TCommandSql::RunCommand(TConfig& config) {
         // Execute query
         settings.ExecMode(NQuery::EExecMode::Execute);
         auto defaultStatsMode = ExplainAnalyzeMode ? NQuery::EStatsMode::Full : NQuery::EStatsMode::None;
-        settings.StatsMode(ParseQueryStatsModeOrThrow(CollectStatsMode, defaultStatsMode));
+        auto statsMode = ParseQueryStatsModeOrThrow(CollectStatsMode, defaultStatsMode);
+        settings.StatsMode(statsMode);
+        if (Progress == "tty") {
+            if (statsMode == NQuery::EStatsMode::None) {
+                throw TMisuseException() << "Non-none statistics collection mode are required to print progress.";
+            }
+            if (statsMode >= NQuery::EStatsMode::Full) {
+                settings.StatsCollectPeriod(std::chrono::milliseconds(3000));
+            } else {
+                settings.StatsCollectPeriod(std::chrono::milliseconds(500));
+            }
+        }
     }
 
     settings.Syntax(SyntaxType);
@@ -190,28 +218,41 @@ int TCommandSql::PrintResponse(NQuery::TExecuteQueryIterator& result) {
     {
         TResultSetPrinter printer(OutputFormat, &IsInterrupted);
 
+        TProgressIndication progressIndication;
+        TMaybe<NQuery::TExecStats> execStats;
+
         while (!IsInterrupted()) {
-            auto streamPart = result.ReadNext().GetValueSync();
+            auto streamPart = result.ReadNext().ExtractValueSync();
             if (ThrowOnErrorAndCheckEOS(streamPart)) {
                 break;
             }
 
+            if (streamPart.HasStats()) {
+                execStats = streamPart.ExtractStats();
+
+                const auto& protoStats = TProtoAccessor::GetProto(execStats.GetRef());
+                for (const auto& queryPhase : protoStats.query_phases()) {
+                    for (const auto& tableAccessStats : queryPhase.table_access()) {
+                        progressIndication.UpdateProgress({tableAccessStats.reads().rows(), tableAccessStats.reads().bytes()});
+                    }
+                }
+                progressIndication.SetDurationUs(protoStats.total_duration_us());
+
+
+                progressIndication.Render();
+            }
+
             if (streamPart.HasResultSet() && !ExplainAnalyzeMode) {
+                progressIndication.Finish();
                 printer.Print(streamPart.GetResultSet());
             }
+        }
 
-            if (streamPart.GetStats().has_value()) {
-                const auto& queryStats = *streamPart.GetStats();
-                stats = queryStats.ToString();
-                ast = queryStats.GetAst();
-
-                if (queryStats.GetPlan()) {
-                    plan = queryStats.GetPlan();
-                }
-                if (queryStats.GetMeta()) {
-                    meta = queryStats.GetMeta();
-                }
-            }
+        if (execStats) {
+            stats = execStats->ToString();
+            plan = execStats->GetPlan();
+            ast = execStats->GetAst();
+            meta = execStats->GetMeta();
         }
     } // TResultSetPrinter destructor should be called before printing stats
 

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

@@ -34,6 +34,7 @@ class TCommandSql : public TYdbCommand, public TCommandWithOutput, public TComma
     bool ExplainMode = false;
     bool ExplainAnalyzeMode = false;
     bool ExplainAst = false;
+    TString Progress;
 };
 
 }

ydb/public/lib/ydb_cli/common/print_utils.cpp

@@ -54,7 +54,7 @@ void PrintSchemeEntry(IOutputStream& o, const NScheme::TSchemeEntry& entry, NCol
 
 TString PrettySize(ui64 size) {
     double sizeFormat = size;
-    return ToString(HumanReadableSize(sizeFormat, ESizeFormat::SF_QUANTITY)) + " B";
+    return ToString(HumanReadableSize(sizeFormat, ESizeFormat::SF_QUANTITY)) + "B";
 }
 
 TString PrettyNumber(ui64 number) {

ydb/public/lib/ydb_cli/common/progress_indication.cpp

@@ -1,21 +1,23 @@
 #include "print_utils.h"
 #include "progress_indication.h"
 
+#include <library/cpp/colorizer/colors.h>
+#include <ydb/public/lib/ydb_cli/common/interactive.h>
+
 namespace NYdb {
 namespace NConsoleClient {
 
-TProgressIndication::TProgressIndication(bool onlyReadStats)
-    : OnlyReadStats(onlyReadStats) {
+TProgressIndication::TProgressIndication() {
 }
 
 void TProgressIndication::UpdateProgress(const TCurrentStats& stats)
 {
     CurrentStats.ReadRows += stats.ReadRows;
     CurrentStats.ReadBytes += stats.ReadBytes;
-    CurrentStats.UpdateRows += stats.UpdateRows;
-    CurrentStats.UpdateBytes += stats.UpdateBytes;
-    CurrentStats.DeleteRows += stats.DeleteRows;
-    CurrentStats.DeleteBytes += stats.DeleteBytes;
+}
+
+void TProgressIndication::SetDurationUs(ui64 durationUs) {
+    DurationUs = durationUs;
 }
 
 TProgressIndication::~TProgressIndication() {
@@ -31,24 +33,31 @@ void TProgressIndication::Finish() {
 
 void TProgressIndication::Render()
 {
+    if (!GetTerminalWidth())
+        return;
+
+    NColorizer::TColors colors = NColorizer::AutoColors(Cout);
+
     Cerr << "\r" "\033[K";
+    Cerr << colors.LightGreen();
     switch (RendersCount % 4) {
         case 0: Cerr << "|"; break;
         case 1: Cerr << "/"; break;
         case 2: Cerr << "-"; break;
         case 3: Cerr << "\\"; break;
     }
-    Cerr << "Progress: ";
-
-    Cerr << PrettyNumber(CurrentStats.ReadRows) << " rows read, " << PrettySize(CurrentStats.ReadBytes) << " read"; 
+    Cerr << colors.Default() << "Progress: " << colors.Default();
 
-    if (OnlyReadStats) {
-        Cerr << ".";
-    } else {
-        Cerr << ", " << PrettyNumber(CurrentStats.UpdateRows) << " rows updated, " << PrettySize(CurrentStats.UpdateBytes) << " updated, " <<
-        PrettyNumber(CurrentStats.DeleteRows) << " rows deleted, " << PrettySize(CurrentStats.DeleteBytes) << " deleted.";
+    Cerr << colors.Default() << PrettyNumber(CurrentStats.ReadRows) << " rows read, " << PrettySize(CurrentStats.ReadBytes) << " read";
+    if (DurationUs) {
+        Cerr << colors.Default() << " (" << PrettyNumber(CurrentStats.ReadRows * 1000000.0 / DurationUs) << "/s" << ", " <<
+            PrettySize(CurrentStats.ReadBytes * 1000000.0 / DurationUs) << "/s" << ")";
     }
 
+    Cerr << colors.Default() << ".";
+
+    Cerr.Flush();
+
     RendersCount++;
 
     CurrentStats = TCurrentStats();

ydb/public/lib/ydb_cli/common/progress_indication.h

@@ -11,16 +11,13 @@ class TProgressIndication {
     struct TCurrentStats {
         ui64 ReadRows = 0;
         ui64 ReadBytes = 0;
-        ui64 UpdateRows = 0;
-        ui64 UpdateBytes = 0;
-        ui64 DeleteRows = 0;
-        ui64 DeleteBytes = 0;
     };
-    explicit TProgressIndication(bool onlyReadStats = false);
+    explicit TProgressIndication();
 
     ~TProgressIndication();
 
     void UpdateProgress(const TCurrentStats& stats);
+    void SetDurationUs(ui64 durationUs);
     void Render();
     void Finish();
 
@@ -29,7 +26,7 @@ class TProgressIndication {
     TCurrentStats CurrentStats;
     bool Finished = false;
     ui32 RendersCount = 0;
-    bool OnlyReadStats = false;
+    ui64 DurationUs = 0;
 };
 
 } // namespace NConsoleClient