implement single part task

Author: arthurpassos

src/Core/Settings.cpp

@@ -6877,8 +6877,8 @@ Experimental timeSeries* aggregate functions for Prometheus-like timeseries resa
     DECLARE_WITH_ALIAS(Bool, allow_experimental_export_merge_tree_partition, false, R"(
 Experimental export merge tree partition.
 )", EXPERIMENTAL, allow_experimental_export_merge_tree_partition) \
-    DECLARE_WITH_ALIAS(Bool, export_merge_tree_partition_background_execution, true, R"(
-    Process exports asynchronously in background threads
+    DECLARE_WITH_ALIAS(Bool, export_merge_tree_partition_individual_part_executor, true, R"(
+    Use the part task instead of the partition task
 )", EXPERIMENTAL, export_merge_tree_partition_background_execution) \
     \
     /* ####################################################### */ \

src/Storages/IStorage.h

@@ -462,6 +462,13 @@ class IStorage : public std::enable_shared_from_this<IStorage>, public TypePromo
         ContextPtr /*context*/,
         std::function<void(MergeTreePartImportStats)>) {}
 
+    virtual void importMergeTreePart(
+          const MergeTreeData &,
+          const DataPartPtr &,
+          ContextPtr,
+          std::function<void(MergeTreePartImportStats)> 
+      ) {}
+
     /** Writes the data to a table in distributed manner.
       * It is supposed that implementation looks into SELECT part of the query and executes distributed
       * INSERT SELECT if it is possible with current storage as a receiver and query SELECT part as a producer.

src/Storages/MergeTree/MergeTreeData.h

@@ -1233,6 +1233,7 @@ class MergeTreeData : public IStorage, public WithMutableContext
     friend class IPartMetadataManager;
     friend class IMergedBlockOutputStream; // for access to log
     friend class ExportPartitionPlainMergeTreeTask;
+    friend class ExportPartPlainMergeTreeTask;
 
     bool require_part_metadata;
 

src/Storages/MergeTree/MergeTreeExportManifest.h

@@ -138,6 +138,11 @@ struct MergeTreeExportManifest
         out->sync();
     }
 
+    void deleteFile()
+    {
+        disk->removeFile(file_path);
+    }
+
     void updateRemotePathAndWrite(const String & part_name, const String & remote_path)
     {
         for (auto & i : items)

src/Storages/ObjectStorage/MergeTree/ExportPartPlainMergeTreeTask.cpp

@@ -0,0 +1,198 @@
+#include <Storages/ObjectStorage/MergeTree/ExportPartPlainMergeTreeTask.h>
+#include <Storages/MergeTree/MergeTreeData.h>
+#include <Storages/MergeTree/MergeTreePartInfo.h>
+#include <Interpreters/PartLog.h>
+#include <Common/logger_useful.h>
+#include <Storages/ObjectStorage/MergeTree/StorageObjectStorageMergeTreePartImporterSink.h>
+#include <Storages/StorageMergeTree.h>
+
+namespace DB
+{
+
+ExportPartPlainMergeTreeTask::ExportPartPlainMergeTreeTask(
+    StorageMergeTree & storage_,
+    const std::shared_ptr<CurrentlyExportingPartsTagger> & exports_tagger_,
+    const StoragePtr & destination_storage_,
+    ContextPtr context_,
+    std::shared_ptr<MergeTreeExportManifest> manifest_,
+    IExecutableTask::TaskResultCallback & task_result_callback_,
+    size_t max_retries_)
+    : storage(storage_)
+    , exports_tagger(exports_tagger_)
+    , destination_storage(destination_storage_)
+    , context(std::move(context_))
+    , manifest(std::move(manifest_))
+    , task_result_callback(task_result_callback_)
+    , max_retries(max_retries_)
+{
+    UInt64 transaction_id = std::stoull(manifest->transaction_id);
+    priority.value = transaction_id;
+}
+
+StorageID ExportPartPlainMergeTreeTask::getStorageID() const
+{
+    return storage.getStorageID();
+}
+
+String ExportPartPlainMergeTreeTask::getQueryId() const
+{
+    return getStorageID().getShortName() + "::export_partition::" + manifest->transaction_id;
+}
+
+bool ExportPartPlainMergeTreeTask::executeStep()
+{
+    if (cancelled)
+        return false;
+
+    switch (state)
+    {
+        case State::NEED_PREPARE:
+        {
+            prepare();
+            state = State::NEED_EXECUTE;
+            return true;
+        }
+        case State::NEED_EXECUTE:
+        {
+            if (executeExport())
+            {
+                state = State::NEED_COMMIT;
+                return true;
+            }
+
+            if (retry_count < max_retries)
+            {
+                retry_count++;
+                LOG_INFO(getLogger("ExportPartPlainMergeTreeTask"),
+                "Retrying export attempt {} for part {}",
+                retry_count, exports_tagger->parts_to_export[0]->name);
+                state = State::NEED_EXECUTE;
+
+                return true;
+            }
+
+            return false;
+        }
+        case State::NEED_COMMIT:
+        {
+            if (commitExport())
+            {
+                state = State::SUCCESS;
+                return true;
+            }
+
+            if (retry_count < max_retries)
+            {
+                retry_count++;
+                LOG_INFO(getLogger("ExportPartPlainMergeTreeTask"),
+                "Retrying export attempt {} for part {}",
+                retry_count, exports_tagger->parts_to_export[0]->name);
+                state = State::NEED_COMMIT;
+
+                return true;
+            }
+
+            return false;
+        }
+        case State::SUCCESS:
+        {
+            return false;
+        }
+    }
+
+    return false;
+}
+
+
+void ExportPartPlainMergeTreeTask::prepare()
+{
+    stopwatch_ptr = std::make_unique<Stopwatch>();
+}
+
+bool ExportPartPlainMergeTreeTask::executeExport()
+{
+    if (cancelled)
+        return false;
+
+    std::function<void(MergeTreePartImportStats)> part_log_wrapper = [this](MergeTreePartImportStats stats) {
+
+        std::lock_guard lock(storage.export_partition_transaction_id_to_manifest_mutex);
+        auto table_id = storage.getStorageID();
+
+        if (stats.status.code != 0)
+        {
+            LOG_INFO(getLogger("ExportMergeTreePartitionToObjectStorageTask"), "Error importing part {}: {}", stats.part->name, stats.status.message);
+            return;
+        }
+
+        storage.export_partition_transaction_id_to_manifest[manifest->transaction_id]->updateRemotePathAndWrite(
+            stats.part->name, 
+            stats.file_path);
+
+        UInt64 elapsed_ns = stopwatch_ptr->elapsedNanoseconds();
+
+        storage.writePartLog(
+            PartLogElement::Type::EXPORT_PART,
+            stats.status,
+            elapsed_ns,
+            stats.part->name,
+            stats.part,
+            {stats.part},
+            nullptr,
+            nullptr);
+    };
+
+    try
+    {
+        destination_storage->importMergeTreePart(
+            storage,
+            exports_tagger->parts_to_export[0],
+            context,
+            part_log_wrapper);
+
+        return true;
+    }
+    catch (...)
+    {
+        LOG_ERROR(getLogger("ExportPartPlainMergeTreeTask"), "Failed to export part {}", exports_tagger->parts_to_export[0]->name);
+        
+        return false;
+    }
+}
+
+bool ExportPartPlainMergeTreeTask::commitExport()
+{
+    std::lock_guard lock(storage.export_partition_transaction_id_to_manifest_mutex);
+
+    if (manifest->exportedPaths().size() == manifest->items.size())
+    {
+        destination_storage->commitExportPartitionTransaction(
+            manifest->transaction_id,
+            manifest->partition_id,
+            manifest->exportedPaths(),
+            context);
+        manifest->completed = true;
+        manifest->write();
+        storage.export_partition_transaction_id_to_manifest.erase(manifest->transaction_id);
+        LOG_INFO(getLogger("ExportMergeTreePartitionToObjectStorageTask"),
+        "Successfully committed export transaction {} for partition {}",
+        manifest->transaction_id, manifest->partition_id);
+    }
+
+    LOG_INFO(getLogger("ExportPartPlainMergeTreeTask"), "Not all parts have been exported yet for transaction id {}, not comitting for this part", manifest->transaction_id);
+
+    return true;
+}
+
+void ExportPartPlainMergeTreeTask::onCompleted()
+{
+    bool success = (state == State::SUCCESS);
+    task_result_callback(success);
+}
+
+void ExportPartPlainMergeTreeTask::cancel() noexcept
+{
+    cancelled = true;
+}
+
+}

src/Storages/ObjectStorage/MergeTree/ExportPartPlainMergeTreeTask.h

@@ -0,0 +1,69 @@
+#pragma once
+
+#include <Storages/MergeTree/IExecutableTask.h>
+#include <Storages/MergeTree/IMergeTreeDataPart.h>
+#include <Storages/MergeTree/MergeTreeExportManifest.h>
+#include <Storages/ObjectStorage/StorageObjectStorage.h>
+#include <Interpreters/Context.h>
+#include <Common/Stopwatch.h>
+#include <Storages/MergeTree/MergeMutateSelectedEntry.h>
+
+namespace DB
+{
+
+class MergeTreeData;
+
+class ExportPartPlainMergeTreeTask : public IExecutableTask
+{
+public:
+    ExportPartPlainMergeTreeTask(
+        StorageMergeTree & storage_,
+        const std::shared_ptr<CurrentlyExportingPartsTagger> & exports_tagger_,
+        const StoragePtr & destination_storage_,
+        ContextPtr context_,
+        std::shared_ptr<MergeTreeExportManifest> manifest_,
+        IExecutableTask::TaskResultCallback & task_result_callback_,
+        size_t max_retries_ = 3);
+
+    void onCompleted() override;
+    bool executeStep() override;
+    void cancel() noexcept override;
+    StorageID getStorageID() const override;
+    Priority getPriority() const override { return priority; }
+    String getQueryId() const override;
+
+private:
+    void prepare();
+    bool executeExport();
+    bool commitExport();
+    bool exportedAllIndividualParts() const;
+
+    enum class State : uint8_t
+    {
+        NEED_PREPARE,
+        NEED_EXECUTE,
+        NEED_COMMIT,
+        SUCCESS
+    };
+
+    State state{State::NEED_PREPARE};
+
+    StorageMergeTree & storage;
+    std::shared_ptr<CurrentlyExportingPartsTagger> exports_tagger;
+    StoragePtr destination_storage;
+    ContextPtr context;
+    std::shared_ptr<MergeTreeExportManifest> manifest;
+    IExecutableTask::TaskResultCallback task_result_callback;
+
+    size_t max_retries;
+    size_t retry_count = 0;
+    Priority priority;
+    std::unique_ptr<Stopwatch> stopwatch_ptr;
+
+    bool cancelled = false;
+    std::exception_ptr current_exception;
+};
+
+using ExportPartPlainMergeTreeTaskPtr = std::shared_ptr<ExportPartPlainMergeTreeTask>;
+
+}

src/Storages/ObjectStorage/MergeTree/ExportPartitionPlainMergeTreeTask.cpp

@@ -168,7 +168,13 @@ void ExportPartitionPlainMergeTreeTask::executeExport()
     {
         LOG_ERROR(getLogger("ExportMergeTreePartitionToObjectStorageTask"),
             "Export attempt failed completely: {}", getCurrentExceptionMessage(true));
-        
+
+        std::lock_guard lock(storage.export_partition_transaction_id_to_manifest_mutex);
+        manifest->deleteFile();
+        storage.already_exported_partition_ids.erase(manifest->partition_id);
+        storage.export_partition_transaction_id_to_manifest.erase(manifest->transaction_id);
+
+
         throw;
     }
 }