introduce file already exists policy setting

Author: arthurpassos

src/Core/Settings.cpp

@@ -6883,6 +6883,12 @@ Maximum number of retries for exporting a merge tree part in an export partition
     DECLARE(UInt64, export_merge_tree_partition_manifest_ttl, 180, R"(
 Determines how long the manifest will live in ZooKeeper. It prevents the same partition from being exported twice to the same destination.
 This setting does not affect / delete in progress tasks. It'll only cleanup the completed ones.
+)", 0) \
+    DECLARE(String, export_merge_tree_part_file_already_exists_policy, "NO_OP", R"(
+Possible values:
+- NO_OP - No-op if the file already exists - Default.
+- ERROR - Throw an error if the file already exists.
+- OVERWRITE - Overwrite the file
 )", 0) \
     \
     /* ####################################################### */ \

src/Storages/ExportReplicatedMergeTreePartitionManifest.h

@@ -5,6 +5,7 @@
 #include <Poco/JSON/Object.h>
 #include <Poco/JSON/Array.h>
 #include <Poco/JSON/Parser.h>
+#include <Storages/MergeTree/MergeTreePartExportManifest.h>
 
 namespace DB
 {
@@ -102,6 +103,7 @@ struct ExportReplicatedMergeTreePartitionManifest
     size_t max_threads;
     bool parallel_formatting;
     bool parquet_parallel_encoding;
+    MergeTreePartExportManifest::FileAlreadyExistsPolicy file_already_exists_policy;
 
     std::string toJsonString() const
     {
@@ -120,6 +122,7 @@ struct ExportReplicatedMergeTreePartitionManifest
         json.set("parallel_formatting", parallel_formatting);
         json.set("max_threads", max_threads);
         json.set("parquet_parallel_encoding", parquet_parallel_encoding);
+        json.set("file_already_exists_policy", String(magic_enum::enum_name(file_already_exists_policy)));
         json.set("create_time", create_time);
         json.set("max_retries", max_retries);
         json.set("ttl_seconds", ttl_seconds);
@@ -152,6 +155,18 @@ struct ExportReplicatedMergeTreePartitionManifest
         manifest.max_threads = json->getValue<size_t>("max_threads");
         manifest.parallel_formatting = json->getValue<bool>("parallel_formatting");
         manifest.parquet_parallel_encoding = json->getValue<bool>("parquet_parallel_encoding");
+
+        if (json->has("file_already_exists_policy"))
+        {
+            const auto file_already_exists_policy = magic_enum::enum_cast<MergeTreePartExportManifest::FileAlreadyExistsPolicy>(json->getValue<String>("file_already_exists_policy"));
+            if (file_already_exists_policy)
+            {
+                manifest.file_already_exists_policy = file_already_exists_policy.value();
+            }
+
+            /// what to do if it's not a valid value?
+        }
+
         return manifest;
     }
 };

src/Storages/MergeTree/ExportList.h

@@ -41,7 +41,7 @@ struct ExportsListElement : private boost::noncopyable
     const StorageID destination_table_id;
     const UInt64 part_size;
     const String part_name;
-    const String destination_file_path;
+    String destination_file_path;
     UInt64 rows_read {0};
     UInt64 total_rows_to_read {0};
     UInt64 total_size_bytes_compressed {0};

src/Storages/MergeTree/ExportPartTask.cpp

@@ -12,6 +12,7 @@
 #include <QueryPipeline/QueryPipelineBuilder.h>
 #include "Common/Exception.h"
 #include <Common/ProfileEventsScope.h>
+#include "Storages/MergeTree/ExportList.h"
 
 namespace ProfileEvents
 {
@@ -68,8 +69,19 @@ bool ExportPartTask::executeStep()
         throw Exception(ErrorCodes::UNKNOWN_TABLE, "Failed to reconstruct destination storage: {}", destination_storage_id_name);
     }
 
+    auto exports_list_entry = storage.getContext()->getExportsList().insert(
+        getStorageID(),
+        manifest.destination_storage_id,
+        manifest.data_part->getBytesOnDisk(),
+        manifest.data_part->name,
+        "not_computed_yet",
+        manifest.data_part->rows_count,
+        manifest.data_part->getBytesOnDisk(),
+        manifest.data_part->getBytesUncompressedOnDisk(),
+        manifest.create_time,
+        local_context);
+
     SinkToStoragePtr sink;
-    std::string destination_file_path;
 
     try
     {
@@ -81,18 +93,44 @@ bool ExportPartTask::executeStep()
         sink = destination_storage->import(
             manifest.data_part->name + "_" + manifest.data_part->checksums.getTotalChecksumHex(),
             block_with_partition_values,
-            destination_file_path,
-            manifest.overwrite_file_if_exists,
+            (*exports_list_entry)->destination_file_path,
+            manifest.file_already_exists_policy == MergeTreePartExportManifest::FileAlreadyExistsPolicy::OVERWRITE,
             context_copy);
     }
     catch (const Exception & e)
     {
-        tryLogCurrentException(__PRETTY_FUNCTION__);
         if (e.code() == ErrorCodes::FILE_ALREADY_EXISTS)
         {
             ProfileEvents::increment(ProfileEvents::PartsExportDuplicated);
+
+            /// File already exists and the policy is NO_OP, treat it as success.
+            if (manifest.file_already_exists_policy == MergeTreePartExportManifest::FileAlreadyExistsPolicy::NO_OP)
+            {
+                storage.writePartLog(
+                    PartLogElement::Type::EXPORT_PART,
+                    {},
+                    static_cast<UInt64>((*exports_list_entry)->elapsed * 1000000000),
+                    manifest.data_part->name,
+                    manifest.data_part,
+                    {manifest.data_part},
+                    nullptr,
+                    nullptr,
+                    exports_list_entry.get());
+
+                std::lock_guard inner_lock(storage.export_manifests_mutex);
+                storage.export_manifests.erase(manifest);
+
+                ProfileEvents::increment(ProfileEvents::PartsExports);
+                ProfileEvents::increment(ProfileEvents::PartsExportTotalMilliseconds, static_cast<UInt64>((*exports_list_entry)->elapsed * 1000));
+
+                if (manifest.completion_callback)
+                    manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createSuccess((*exports_list_entry)->destination_file_path));
+                return false;
+            }
         }
 
+        tryLogCurrentException(__PRETTY_FUNCTION__);
+
         ProfileEvents::increment(ProfileEvents::PartsExportFailures);
 
         std::lock_guard inner_lock(storage.export_manifests_mutex);
@@ -139,17 +177,6 @@ bool ExportPartTask::executeStep()
         local_context,
         getLogger("ExportPartition"));
 
-    auto exports_list_entry = storage.getContext()->getExportsList().insert(
-        getStorageID(),
-        manifest.destination_storage_id,
-        manifest.data_part->getBytesOnDisk(),
-        manifest.data_part->name,
-        destination_file_path,
-        manifest.data_part->rows_count,
-        manifest.data_part->getBytesOnDisk(),
-        manifest.data_part->getBytesUncompressedOnDisk(),
-        manifest.create_time,
-        local_context);
 
     ThreadGroupSwitcher switcher((*exports_list_entry)->thread_group, "");
 
@@ -198,7 +225,7 @@ bool ExportPartTask::executeStep()
         ProfileEvents::increment(ProfileEvents::PartsExportTotalMilliseconds, static_cast<UInt64>((*exports_list_entry)->elapsed * 1000));
 
         if (manifest.completion_callback)
-            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createSuccess(destination_file_path));
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createSuccess((*exports_list_entry)->destination_file_path));
     }
     catch (const Exception & e)
     {

src/Storages/MergeTree/ExportPartitionTaskScheduler.cpp

@@ -4,6 +4,7 @@
 #include <Interpreters/DatabaseCatalog.h>
 #include "Common/ZooKeeper/Types.h"
 #include "Storages/MergeTree/ExportPartitionUtils.h"
+#include "Storages/MergeTree/MergeTreePartExportManifest.h"
 
 
 namespace DB
@@ -17,6 +18,7 @@ namespace
         context_copy->setSetting("output_format_parallel_formatting", manifest.parallel_formatting);
         context_copy->setSetting("output_format_parquet_parallel_encoding", manifest.parquet_parallel_encoding);
         context_copy->setSetting("max_threads", manifest.max_threads);
+        context_copy->setSetting("export_merge_tree_part_file_already_exists_policy", String(magic_enum::enum_name(manifest.file_already_exists_policy)));
         return context_copy;
     }
 }

src/Storages/MergeTree/MergeTreeData.cpp

@@ -212,7 +212,7 @@ namespace Setting
     extern const SettingsBool apply_patch_parts;
     extern const SettingsBool allow_experimental_export_merge_tree_part;
     extern const SettingsUInt64 min_bytes_to_use_direct_io;
-    extern const SettingsBool export_merge_tree_part_overwrite_file_if_exists;
+    extern const SettingsString export_merge_tree_part_file_already_exists_policy;
     extern const SettingsBool output_format_parallel_formatting;
     extern const SettingsBool output_format_parquet_parallel_encoding;
 }
@@ -6253,11 +6253,20 @@ void MergeTreeData::exportPartToTable(
                         part_name, getStorageID().getFullTableName());
 
     {
+        const auto file_already_exists_policy = magic_enum::enum_cast<MergeTreePartExportManifest::FileAlreadyExistsPolicy>(Poco::toUpper(query_context->getSettingsRef()[Setting::export_merge_tree_part_file_already_exists_policy].value));
+        if (!file_already_exists_policy)
+        {
+            throw Exception(
+                ErrorCodes::INVALID_SETTING_VALUE,
+                "Invalid value for setting 'export_merge_tree_part_file_already_exists_policy': {}",
+                query_context->getSettingsRef()[Setting::export_merge_tree_part_file_already_exists_policy].value);
+        }
+
         MergeTreePartExportManifest manifest(
             dest_storage->getStorageID(),
             part,
             transaction_id,
-            query_context->getSettingsRef()[Setting::export_merge_tree_part_overwrite_file_if_exists],
+            *file_already_exists_policy,
             query_context->getSettingsRef()[Setting::output_format_parallel_formatting],
             query_context->getSettingsRef()[Setting::output_format_parquet_parallel_encoding],
             query_context->getSettingsRef()[Setting::max_threads],

src/Storages/MergeTree/MergeTreePartExportManifest.h

@@ -11,6 +11,13 @@ class ExportPartTask;
 
 struct MergeTreePartExportManifest
 {
+    enum class FileAlreadyExistsPolicy
+    {
+        NO_OP,
+        ERROR,
+        OVERWRITE,
+    };
+
     using DataPartPtr = std::shared_ptr<const IMergeTreeDataPart>;
 
     struct CompletionCallbackResult
@@ -39,15 +46,15 @@ struct MergeTreePartExportManifest
         const StorageID & destination_storage_id_,
         const DataPartPtr & data_part_,
         const String & query_id_,
-        bool overwrite_file_if_exists_,
+        FileAlreadyExistsPolicy file_already_exists_policy_,
         bool parallel_formatting_,
         bool parquet_parallel_encoding_,
         std::size_t max_threads_,
         std::function<void(CompletionCallbackResult)> completion_callback_ = {})
         : destination_storage_id(destination_storage_id_),
           data_part(data_part_),
           query_id(query_id_),
-          overwrite_file_if_exists(overwrite_file_if_exists_),
+          file_already_exists_policy(file_already_exists_policy_),
           parallel_formatting(parallel_formatting_),
           parquet_parallel_encoding(parquet_parallel_encoding_),
           max_threads(max_threads_),
@@ -58,7 +65,7 @@ struct MergeTreePartExportManifest
     DataPartPtr data_part;
     /// Used for killing the export.
     String query_id;
-    bool overwrite_file_if_exists;
+    FileAlreadyExistsPolicy file_already_exists_policy;
     bool parallel_formatting;
     /// parquet has a different setting for parallel formatting
     bool parquet_parallel_encoding;

src/Storages/StorageReplicatedMergeTree.cpp

@@ -198,6 +198,7 @@ namespace Setting
     extern const SettingsBool output_format_parallel_formatting;
     extern const SettingsBool output_format_parquet_parallel_encoding;
     extern const SettingsMaxThreads max_threads;
+    extern const SettingsString export_merge_tree_part_file_already_exists_policy;
 }
 
 namespace MergeTreeSetting
@@ -300,6 +301,7 @@ namespace ErrorCodes
     extern const int FAULT_INJECTED;
     extern const int CANNOT_FORGET_PARTITION;
     extern const int TIMEOUT_EXCEEDED;
+    extern const int INVALID_SETTING_VALUE;
 }
 
 namespace ActionLocks
@@ -8159,6 +8161,18 @@ void StorageReplicatedMergeTree::exportPartitionToTable(const PartitionCommand &
     manifest.parallel_formatting = query_context->getSettingsRef()[Setting::output_format_parallel_formatting];
     manifest.parquet_parallel_encoding = query_context->getSettingsRef()[Setting::output_format_parquet_parallel_encoding];
 
+    const auto file_already_exists_policy = magic_enum::enum_cast<MergeTreePartExportManifest::FileAlreadyExistsPolicy>(query_context->getSettingsRef()[Setting::export_merge_tree_part_file_already_exists_policy].value);
+
+    if (!file_already_exists_policy)
+    {
+        throw Exception(
+            ErrorCodes::INVALID_SETTING_VALUE,
+            "Invalid value for setting 'export_merge_tree_part_file_already_exists_policy': {}",
+            query_context->getSettingsRef()[Setting::export_merge_tree_part_file_already_exists_policy].value);
+    }
+
+    manifest.file_already_exists_policy = file_already_exists_policy.value();
+
     ops.emplace_back(zkutil::makeCreateRequest(
         fs::path(partition_exports_path) / "metadata.json", 
         manifest.toJsonString(),