Fix progress bar for local

programs/local/LocalServer.cpp

@@ -388,24 +388,32 @@ void LocalServer::processQueries()
     /// Use the same query_id (and thread group) for all queries
     CurrentThread::QueryScope query_scope_holder(context);
 
-    ///Set progress show
+    /// Set progress show
     need_render_progress = config().getBool("progress", false);
 
+    std::function<void()> finalize_progress;
     if (need_render_progress)
     {
+        /// Set progress callback, which can be run from multiple threads.
         context->setProgressCallback([&](const Progress & value)
         {
             /// Write progress only if progress was updated
             if (progress_indication.updateProgress(value))
                 progress_indication.writeProgress();
         });
+
+        /// Set finalizing callback for progress, which is called right before finalizing query output.
+        finalize_progress = [&]()
+        {
+            progress_indication.clearProgressOutput();
+        };
+
+        /// Set callback for file processing progress.
+        progress_indication.setFileProgressCallback(context);
     }
 
     bool echo_queries = config().hasOption("echo") || config().hasOption("verbose");
 
-    if (need_render_progress)
-        progress_indication.setFileProgressCallback(context);
-
     std::exception_ptr exception;
 
     for (const auto & query : queries)
@@ -425,7 +433,7 @@ void LocalServer::processQueries()
 
         try
         {
-            executeQuery(read_buf, write_buf, /* allow_into_outfile = */ true, context, {});
+            executeQuery(read_buf, write_buf, /* allow_into_outfile = */ true, context, {}, finalize_progress);
         }
         catch (...)
         {

src/Common/ProgressIndication.cpp

@@ -4,9 +4,6 @@
 #include <Common/UnicodeBar.h>
 #include <Databases/DatabaseMemory.h>
 
-/// FIXME: progress bar in clickhouse-local needs to be cleared after query execution
-///        - same as it is now in clickhouse-client. Also there is no writeFinalProgress call
-///        in clickhouse-local.
 
 namespace DB
 {

src/Interpreters/executeQuery.cpp

@@ -948,7 +948,8 @@ void executeQuery(
     WriteBuffer & ostr,
     bool allow_into_outfile,
     ContextMutablePtr context,
-    std::function<void(const String &, const String &, const String &, const String &)> set_result_details)
+    std::function<void(const String &, const String &, const String &, const String &)> set_result_details,
+    std::function<void()> before_finalize_callback)
 {
     PODArray<char> parse_buf;
     const char * begin;
@@ -1079,6 +1080,8 @@ void executeQuery(
                     out->onProgress(progress);
                 });
 
+                out->setBeforeFinalizeCallback(before_finalize_callback);
+
                 if (set_result_details)
                     set_result_details(
                         context->getClientInfo().current_query_id, out->getContentType(), format_name, DateLUT::instance().getTimeZone());

src/Interpreters/executeQuery.h

@@ -17,7 +17,8 @@ void executeQuery(
     WriteBuffer & ostr,                 /// Where to write query output to.
     bool allow_into_outfile,            /// If true and the query contains INTO OUTFILE section, redirect output to that file.
     ContextMutablePtr context,                 /// DB, tables, data types, storage engines, functions, aggregate functions...
-    std::function<void(const String &, const String &, const String &, const String &)> set_result_details /// If a non-empty callback is passed, it will be called with the query id, the content-type, the format, and the timezone.
+    std::function<void(const String &, const String &, const String &, const String &)> set_result_details, /// If a non-empty callback is passed, it will be called with the query id, the content-type, the format, and the timezone.
+    std::function<void()> before_finalize_callback = {} /// Will be set in output format to be called before finalize.
 );
 
 

src/Processors/Formats/IOutputFormat.cpp

@@ -76,6 +76,9 @@ void IOutputFormat::work()
         if (rows_before_limit_counter && rows_before_limit_counter->hasAppliedLimit())
             setRowsBeforeLimit(rows_before_limit_counter->get());
 
+        if (before_finalize_callback)
+            before_finalize_callback();
+
         finalize();
         finalized = true;
         return;
@@ -117,4 +120,3 @@ void IOutputFormat::write(const Block & block)
 }
 
 }
-

src/Processors/Formats/IOutputFormat.h

@@ -67,6 +67,9 @@ public:
     /// Passed value are delta, that must be summarized.
     virtual void onProgress(const Progress & /*progress*/) {}
 
+    /// Set callback, which will be called before call to finalize().
+    void setBeforeFinalizeCallback(std::function<void()> callback) { before_finalize_callback = callback; }
+
     /// Content-Type to set when sending HTTP response.
     virtual std::string getContentType() const { return "text/plain; charset=UTF-8"; }
 
@@ -91,6 +94,7 @@ private:
     size_t result_bytes = 0;
 
     bool prefix_written = false;
+
+    std::function<void()> before_finalize_callback;
 };
 }
-