Add processlist to PipelineExecutor. Throw exception if query was killed.

dbms/src/Processors/Executors/PipelineExecutor.cpp

@@ -10,6 +10,7 @@
 #include <Common/Stopwatch.h>
 #include <Processors/ISource.h>
 #include <Common/setThreadName.h>
+#include <Interpreters/ProcessList.h>
 
 #if !defined(__APPLE__) && !defined(__FreeBSD__)
 #include <sched.h>
@@ -33,12 +34,13 @@ static bool checkCanAddAdditionalInfoToException(const DB::Exception & exception
            && exception.code() != ErrorCodes::QUERY_WAS_CANCELLED;
 }
 
-PipelineExecutor::PipelineExecutor(Processors & processors_)
+PipelineExecutor::PipelineExecutor(Processors & processors_, QueryStatus * elem)
     : processors(processors_)
     , cancelled(false)
     , finished(false)
     , num_processing_executors(0)
     , expand_pipeline_task(nullptr)
+    , process_list_element(elem)
 {
     buildGraph();
 }
@@ -473,7 +475,12 @@ void PipelineExecutor::execute(size_t num_threads)
     }
 
     if (cancelled)
+    {
+        if (process_list_element && process_list_element->isKilled())
+            throw Exception("Query was cancelled", ErrorCodes::QUERY_WAS_CANCELLED);
+
         return;
+    }
 
     bool all_processors_finished = true;
     for (auto & node : graph)

dbms/src/Processors/Executors/PipelineExecutor.h

@@ -13,6 +13,7 @@
 namespace DB
 {
 
+class QueryStatus;
 
 /// Executes query pipeline.
 class PipelineExecutor
@@ -24,7 +25,7 @@ public:
     /// During pipeline execution new processors can appear. They will be added to existing set.
     ///
     /// Explicit graph representation is built in constructor. Throws if graph is not correct.
-    explicit PipelineExecutor(Processors & processors_);
+    explicit PipelineExecutor(Processors & processors_, QueryStatus * elem = nullptr);
 
     /// Execute pipeline in multiple threads. Must be called once.
     /// In case of exception during execution throws any occurred.
@@ -242,6 +243,9 @@ private:
     using ProcessorsMap = std::unordered_map<const IProcessor *, UInt64>;
     ProcessorsMap processors_map;
 
+    /// Now it's used to check if query was killed.
+    QueryStatus * process_list_element = nullptr;
+
     /// Graph related methods.
     bool addEdges(UInt64 node);
     void buildGraph();

dbms/src/Processors/QueryPipeline.cpp

@@ -523,6 +523,8 @@ void QueryPipeline::setProgressCallback(const ProgressCallback & callback)
 
 void QueryPipeline::setProcessListElement(QueryStatus * elem)
 {
+    process_list_element = elem;
+
     for (auto & processor : processors)
     {
         if (auto * source = dynamic_cast<ISourceWithProgress *>(processor.get()))
@@ -630,7 +632,7 @@ PipelineExecutorPtr QueryPipeline::execute()
     if (!output_format)
         throw Exception("Cannot execute pipeline because it doesn't have output.", ErrorCodes::LOGICAL_ERROR);
 
-    return std::make_shared<PipelineExecutor>(processors);
+    return std::make_shared<PipelineExecutor>(processors, process_list_element);
 }
 
 }

dbms/src/Processors/QueryPipeline.h

@@ -120,6 +120,8 @@ private:
 
     size_t max_threads = 0;
 
+    QueryStatus * process_list_element = nullptr;
+
     void checkInitialized();
     void checkSource(const ProcessorPtr & source, bool can_have_totals);