Add progress bar to LocalServer

programs/local/LocalServer.cpp

@@ -18,7 +18,9 @@
 #include <Common/escapeForFileName.h>
 #include <Common/ClickHouseRevision.h>
 #include <Common/ThreadStatus.h>
+#include <Common/UnicodeBar.h>
 #include <Common/config_version.h>
+#include <Common/escapeForFileName.h>
 #include <Common/quoteString.h>
 #include <IO/ReadBufferFromFile.h>
 #include <IO/ReadBufferFromString.h>
@@ -45,6 +47,9 @@
 #include <filesystem>
 
 
+/// http://en.wikipedia.org/wiki/ANSI_escape_code
+#define CLEAR_TO_END_OF_LINE "\033[K"
+
 namespace DB
 {
 
@@ -387,11 +392,24 @@ void LocalServer::processQueries()
     /// Use the same query_id (and thread group) for all queries
     CurrentThread::QueryScope query_scope_holder(context);
 
+    context.setProgressCallback([this](const Progress &value){
+        this->updateProgress(value);
+        this->writeProgress();
+        return true;
+    });
+
     bool echo_queries = config().hasOption("echo") || config().hasOption("verbose");
     std::exception_ptr exception;
 
     for (const auto & query : queries)
     {
+        watch.restart();
+        progress.reset();
+        show_progress_bar = false;
+        written_progress_chars = 0;
+        written_first_block = false;
+
+
         ReadBufferFromString read_buf(query);
         WriteBufferFromFileDescriptor write_buf(STDOUT_FILENO);
 
@@ -624,6 +642,95 @@ void LocalServer::applyCmdOptions(ContextPtr context)
     applyCmdSettings(context);
 }
 
+void LocalServer::writeProgress()
+{
+    if (!need_render_progress)
+        return;
+
+    /// Output all progress bar commands to stderr at once to avoid flicker.
+    WriteBufferFromFileDescriptor message(STDERR_FILENO, 1024);
+
+    static size_t increment = 0;
+    static const char * indicators[8] = {
+        "\033[1;30m→\033[0m",
+        "\033[1;31m↘\033[0m",
+        "\033[1;32m↓\033[0m",
+        "\033[1;33m↙\033[0m",
+        "\033[1;34m←\033[0m",
+        "\033[1;35m↖\033[0m",
+        "\033[1;36m↑\033[0m",
+        "\033[1m↗\033[0m",
+    };
+
+    const char * indicator = indicators[increment % 8];
+
+    size_t terminal_width = getTerminalWidth();
+
+    if (!written_progress_chars)
+    {
+        /// If the current line is not empty, the progress must be output on the next line.
+        /// The trick is found here: https://www.vidarholen.net/contents/blog/?p=878
+        message << std::string(terminal_width, ' ');
+    }
+    message << '\r';
+
+    size_t prefix_size = message.count();
+
+    message << indicator << " Progress: ";
+
+    message << formatReadableQuantity(progress.read_rows) << " rows, " << formatReadableSizeWithDecimalSuffix(progress.read_bytes);
+
+    size_t elapsed_ns = watch.elapsed();
+    if (elapsed_ns)
+        message << " (" << formatReadableQuantity(progress.read_rows * 1000000000.0 / elapsed_ns) << " rows/s., "
+                << formatReadableSizeWithDecimalSuffix(progress.read_bytes * 1000000000.0 / elapsed_ns) << "/s.) ";
+    else
+        message << ". ";
+
+    written_progress_chars = message.count() - prefix_size - (strlen(indicator) - 2); /// Don't count invisible output (escape sequences).
+
+    /// If the approximate number of rows to process is known, we can display a progress bar and percentage.
+    if (progress.total_rows_to_read > 0)
+    {
+        size_t total_rows_corrected = std::max(progress.read_rows, progress.total_rows_to_read);
+
+        /// To avoid flicker, display progress bar only if .5 seconds have passed since query execution start
+        ///  and the query is less than halfway done.
+
+        if (elapsed_ns > 500000000)
+        {
+            /// Trigger to start displaying progress bar. If query is mostly done, don't display it.
+            if (progress.read_rows * 2 < total_rows_corrected)
+                show_progress_bar = true;
+
+            if (show_progress_bar)
+            {
+                ssize_t width_of_progress_bar = static_cast<ssize_t>(terminal_width) - written_progress_chars - strlen(" 99%");
+                if (width_of_progress_bar > 0)
+                {
+                    std::string bar
+                        = UnicodeBar::render(UnicodeBar::getWidth(progress.read_rows, 0, total_rows_corrected, width_of_progress_bar));
+                    message << "\033[0;32m" << bar << "\033[0m";
+                    if (width_of_progress_bar > static_cast<ssize_t>(bar.size() / UNICODE_BAR_CHAR_SIZE))
+                        message << std::string(width_of_progress_bar - bar.size() / UNICODE_BAR_CHAR_SIZE, ' ');
+                }
+            }
+        }
+
+        /// Underestimate percentage a bit to avoid displaying 100%.
+        message << ' ' << (99 * progress.read_rows / total_rows_corrected) << '%';
+    }
+
+    message << CLEAR_TO_END_OF_LINE;
+    ++increment;
+
+    message.next();
+}
+void LocalServer::updateProgress(const Progress & value)
+{
+    this->progress.incrementPiecewiseAtomically(value);
+}
+
 }
 
 #pragma GCC diagnostic ignored "-Wunused-function"

programs/local/LocalServer.h

@@ -42,6 +42,9 @@ private:
     void setupUsers();
     void cleanup();
 
+    void updateProgress(const Progress &value);
+    void writeProgress();
+
 protected:
     SharedContextHolder shared_context;
     ContextPtr global_context;
@@ -49,6 +52,15 @@ protected:
     /// Settings specified via command line args
     Settings cmd_settings;
 
+    /// The server periodically sends information about how much data was read since last time.
+    bool need_render_progress = true;    /// Render query execution progress.
+    Progress progress;
+    bool show_progress_bar = false;
+    Stopwatch watch;
+
+    size_t written_progress_chars = 0;
+    bool written_first_block = false;
+
     std::optional<std::filesystem::path> temporary_directory_to_delete;
 };