do not mark part export as failed in case of cancelled export

Author: arthurpassos

src/Storages/MergeTree/ExportPartTask.cpp

@@ -138,7 +138,7 @@ bool ExportPartTask::executeStep()
         storage.export_manifests.erase(manifest);
 
         if (manifest.completion_callback)
-            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createFailure(e.message()));
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createFailure(e));
         return false;
     }
 
@@ -255,7 +255,7 @@ bool ExportPartTask::executeStep()
         storage.export_manifests.erase(manifest);
 
         if (manifest.completion_callback)
-            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createFailure(e.message()));
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createFailure(e));
 
         throw;
     }

src/Storages/MergeTree/ExportPartitionTaskScheduler.cpp

@@ -2,14 +2,21 @@
 #include <Storages/StorageReplicatedMergeTree.h>
 #include <Interpreters/Context.h>
 #include <Interpreters/DatabaseCatalog.h>
-#include "Common/ZooKeeper/Types.h"
+#include <Common/Exception.h>
+#include <Common/ZooKeeper/Types.h>
 #include "Storages/MergeTree/ExportPartitionUtils.h"
 #include "Storages/MergeTree/MergeTreePartExportManifest.h"
 
 
 namespace DB
 {
 
+namespace ErrorCodes
+{
+    extern const int QUERY_WAS_CANCELLED;
+    extern const int LOGICAL_ERROR;
+}
+
 namespace
 {
     ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ExportReplicatedMergeTreePartitionManifest & manifest)
@@ -210,10 +217,22 @@ void ExportPartitionTaskScheduler::handlePartExportFailure(
     const std::string & part_name,
     const std::filesystem::path & export_path,
     const zkutil::ZooKeeperPtr & zk,
-    const String & exception,
+    const std::optional<Exception> & exception,
     size_t max_retries
 )
 {
+    if (!exception)
+    {
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "ExportPartition scheduler task: No exception provided for error handling. Sounds like a bug");
+    }
+
+    /// Early exit if the query was cancelled - no need to increment error counts
+    if (exception->code() == ErrorCodes::QUERY_WAS_CANCELLED)
+    {
+        LOG_INFO(storage.log, "ExportPartition scheduler task: Part {} export was cancelled, skipping error handling", part_name);
+        return;
+    }
+
     Coordination::Stat locked_by_stat;
     std::string locked_by;
 
@@ -274,15 +293,15 @@ void ExportPartitionTaskScheduler::handlePartExportFailure(
             num_exceptions = std::stoull(num_exceptions_string.c_str());
 
             ops.emplace_back(zkutil::makeSetRequest(last_exception_path / "part", part_name, -1));
-            ops.emplace_back(zkutil::makeSetRequest(last_exception_path / "exception", exception, -1));
+            ops.emplace_back(zkutil::makeSetRequest(last_exception_path / "exception", exception->message(), -1));
         }
         else
         {
             ops.emplace_back(zkutil::makeCreateRequest(exceptions_per_replica_path, "", zkutil::CreateMode::Persistent));
             ops.emplace_back(zkutil::makeCreateRequest(count_path, "0", zkutil::CreateMode::Persistent));
             ops.emplace_back(zkutil::makeCreateRequest(last_exception_path, "", zkutil::CreateMode::Persistent));
             ops.emplace_back(zkutil::makeCreateRequest(last_exception_path / "part", part_name, zkutil::CreateMode::Persistent));
-            ops.emplace_back(zkutil::makeCreateRequest(last_exception_path / "exception", exception, zkutil::CreateMode::Persistent));
+            ops.emplace_back(zkutil::makeCreateRequest(last_exception_path / "exception", exception->message(), zkutil::CreateMode::Persistent));
         }
 
         num_exceptions++;

src/Storages/MergeTree/ExportPartitionTaskScheduler.h

@@ -6,6 +6,7 @@
 namespace DB
 {
 
+class Exception;
 class StorageReplicatedMergeTree;
 
 struct ExportReplicatedMergeTreePartitionManifest;
@@ -44,7 +45,7 @@ class ExportPartitionTaskScheduler
         const std::string & part_name,
         const std::filesystem::path & export_path,
         const zkutil::ZooKeeperPtr & zk,
-        const String & exception,
+        const std::optional<Exception> & exception,
         size_t max_retries);
 
     bool tryToMovePartToProcessed(

src/Storages/MergeTree/MergeTreePartExportManifest.h

@@ -3,10 +3,13 @@
 #include <Interpreters/StorageID.h>
 #include <Storages/MergeTree/IMergeTreeDataPart.h>
 #include <QueryPipeline/QueryPipeline.h>
+#include <optional>
 
 namespace DB
 {
 
+class Exception;
+
 class ExportPartTask;
 
 struct MergeTreePartExportManifest
@@ -23,23 +26,23 @@ struct MergeTreePartExportManifest
     struct CompletionCallbackResult
     {
     private:
-        CompletionCallbackResult(bool success_, const String & relative_path_in_destination_storage_, const String & exception_)
-            : success(success_), relative_path_in_destination_storage(relative_path_in_destination_storage_), exception(exception_) {}
+        CompletionCallbackResult(bool success_, const String & relative_path_in_destination_storage_, std::optional<Exception> exception_)
+            : success(success_), relative_path_in_destination_storage(relative_path_in_destination_storage_), exception(std::move(exception_)) {}
     public:
 
         static CompletionCallbackResult createSuccess(const String & relative_path_in_destination_storage_)
         {
-            return CompletionCallbackResult(true, relative_path_in_destination_storage_, "");
+            return CompletionCallbackResult(true, relative_path_in_destination_storage_, std::nullopt);
         }
 
-        static CompletionCallbackResult createFailure(const String & exception_)
+        static CompletionCallbackResult createFailure(Exception exception_)
         {
-            return CompletionCallbackResult(false, "", exception_);
+            return CompletionCallbackResult(false, "", std::move(exception_));
         }
 
         bool success = false;
         String relative_path_in_destination_storage;
-        String exception;
+        std::optional<Exception> exception;
     };
 
     MergeTreePartExportManifest(