use enum for statuses, remove status from completed part export, try to make tests a bit more stable, cache status

Author: arthurpassos

src/Core/SettingsChangesHistory.cpp

@@ -50,6 +50,7 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
             {"export_merge_tree_partition_force_export", false, false, "New setting."},
             {"export_merge_tree_partition_max_retries", 3, 3, "New setting."},
             {"export_merge_tree_partition_manifest_ttl", 180, 180, "New setting."},
+            {"export_merge_tree_part_file_already_exists_policy", "NO_OP", "NO_OP", "New setting."},
         });
         addSettingsChanges(settings_changes_history, "25.8",
         {

src/Storages/ExportReplicatedMergeTreePartitionManifest.h

@@ -12,8 +12,16 @@ namespace DB
 
 struct ExportReplicatedMergeTreePartitionProcessingPartEntry
 {
+
+    enum class Status
+    {
+        PENDING,
+        COMPLETED,
+        FAILED
+    };
+
     String part_name;
-    String status;
+    Status status;
     size_t retry_count;
     String finished_by;
 
@@ -22,7 +30,7 @@ struct ExportReplicatedMergeTreePartitionProcessingPartEntry
         Poco::JSON::Object json;
 
         json.set("part_name", part_name);
-        json.set("status", status);
+        json.set("status", String(magic_enum::enum_name(status)));
         json.set("retry_count", retry_count);
         json.set("finished_by", finished_by);
         std::ostringstream oss;     // STYLE_CHECK_ALLOW_STD_STRING_STREAM
@@ -41,7 +49,7 @@ struct ExportReplicatedMergeTreePartitionProcessingPartEntry
         ExportReplicatedMergeTreePartitionProcessingPartEntry entry;
 
         entry.part_name = json->getValue<String>("part_name");
-        entry.status = json->getValue<String>("status");
+        entry.status = magic_enum::enum_cast<Status>(json->getValue<String>("status")).value();
         entry.retry_count = json->getValue<size_t>("retry_count");
         if (json->has("finished_by"))
         {
@@ -55,15 +63,13 @@ struct ExportReplicatedMergeTreePartitionProcessedPartEntry
 {
     String part_name;
     String path_in_destination;
-    String status;
     String finished_by;
 
     std::string toJsonString() const
     {
         Poco::JSON::Object json;
         json.set("part_name", part_name);
         json.set("path_in_destination", path_in_destination);
-        json.set("status", status);
         json.set("finished_by", finished_by);
         std::ostringstream oss;     // STYLE_CHECK_ALLOW_STD_STRING_STREAM
         oss.exceptions(std::ios::failbit);
@@ -81,7 +87,6 @@ struct ExportReplicatedMergeTreePartitionProcessedPartEntry
 
         entry.part_name = json->getValue<String>("part_name");
         entry.path_in_destination = json->getValue<String>("path_in_destination");
-        entry.status = json->getValue<String>("status");
         entry.finished_by = json->getValue<String>("finished_by");
 
         return entry;

src/Storages/ExportReplicatedMergeTreePartitionTaskEntry.h

@@ -15,6 +15,16 @@ struct ExportReplicatedMergeTreePartitionTaskEntry
     using DataPartPtr = std::shared_ptr<const IMergeTreeDataPart>;
     ExportReplicatedMergeTreePartitionManifest manifest;
 
+    enum class Status
+    {
+        PENDING,
+        COMPLETED,
+        FAILED
+    };
+
+    /// Allows us to skip completed / failed entries during scheduling
+    mutable Status status;
+
     /// References to the parts that should be exported
     /// This is used to prevent the parts from being deleted before finishing the export operation
     /// It does not mean this replica will export all the parts

src/Storages/MergeTree/ExportPartitionManifestUpdatingTask.cpp

@@ -265,7 +265,7 @@ void ExportPartitionManifestUpdatingTask::addTask(
     }
 
     /// Insert or update entry. The multi_index container automatically maintains both indexes.
-    auto entry = ExportReplicatedMergeTreePartitionTaskEntry {metadata, std::move(part_references)};
+    auto entry = ExportReplicatedMergeTreePartitionTaskEntry {metadata, ExportReplicatedMergeTreePartitionTaskEntry::Status::PENDING, std::move(part_references)};
     auto it = entries_by_key.find(key);
     if (it != entries_by_key.end())
         entries_by_key.replace(it, entry);

src/Storages/MergeTree/ExportPartitionTaskScheduler.cpp

@@ -35,13 +35,20 @@ void ExportPartitionTaskScheduler::run()
     auto zk = storage.getZooKeeper();
 
     // Iterate sorted by create_time
-    for (const auto & entry : storage.export_merge_tree_partition_task_entries_by_create_time)
+    for (auto & entry : storage.export_merge_tree_partition_task_entries_by_create_time)
     {
         const auto & manifest = entry.manifest;
         const auto key = entry.getCompositeKey();
         const auto & database = storage.getContext()->resolveDatabase(manifest.destination_database);
         const auto & table = manifest.destination_table;
 
+        /// No need to query zk for status if the local one is not PENDING
+        if (entry.status != ExportReplicatedMergeTreePartitionTaskEntry::Status::PENDING)
+        {
+            LOG_INFO(storage.log, "ExportPartition scheduler task: Skipping... Local status is {}", magic_enum::enum_name(entry.status).data());
+            continue;
+        }
+
         const auto destination_storage_id = StorageID(QualifiedTableName {database, table});
 
         const auto destination_storage = DatabaseCatalog::instance().tryGetTable(destination_storage_id, storage.getContext());
@@ -52,16 +59,25 @@ void ExportPartitionTaskScheduler::run()
             continue;
         }
 
-        std::string status;
-        if (!zk->tryGet(fs::path(storage.zookeeper_path) / "exports" / key / "status", status))
+        std::string status_in_zk_string;
+        if (!zk->tryGet(fs::path(storage.zookeeper_path) / "exports" / key / "status", status_in_zk_string))
         {
             LOG_INFO(storage.log, "ExportPartition scheduler task: Failed to get status, skipping");
             continue;
         }
 
-        if (status != "PENDING")
+        const auto status_in_zk = magic_enum::enum_cast<ExportReplicatedMergeTreePartitionTaskEntry::Status>(status_in_zk_string);
+
+        if (!status_in_zk)
+        {
+            LOG_INFO(storage.log, "ExportPartition scheduler task: Failed to get status from zk, skipping");
+            continue;
+        }
+
+        if (status_in_zk.value() != ExportReplicatedMergeTreePartitionTaskEntry::Status::PENDING)
         {
-            LOG_INFO(storage.log, "ExportPartition scheduler task: Skipping... Status is not PENDING");
+            entry.status = status_in_zk.value();
+            LOG_INFO(storage.log, "ExportPartition scheduler task: Skipping... Status from zk is {}", entry.status);
             continue;
         }
 
@@ -238,10 +254,10 @@ void ExportPartitionTaskScheduler::handlePartExportFailure(
         if (processing_part_entry.retry_count >= max_retries)
         {
             /// just set status in processing_part_path and finished_by
-            processing_part_entry.status = "FAILED";
+            processing_part_entry.status = ExportReplicatedMergeTreePartitionProcessingPartEntry::Status::FAILED;
             processing_part_entry.finished_by = storage.replica_name;
 
-            ops.emplace_back(zkutil::makeSetRequest(export_path / "status", "FAILED", -1));
+            ops.emplace_back(zkutil::makeSetRequest(export_path / "status", String(magic_enum::enum_name(ExportReplicatedMergeTreePartitionTaskEntry::Status::FAILED)).data(), -1));
             LOG_INFO(storage.log, "ExportPartition scheduler task: Retry count limit exceeded for part {}, will try to fail the entire task", part_name);
         }
 
@@ -312,7 +328,6 @@ bool ExportPartitionTaskScheduler::tryToMovePartToProcessed(
     ExportReplicatedMergeTreePartitionProcessedPartEntry processed_part_entry;
     processed_part_entry.part_name = part_name;
     processed_part_entry.path_in_destination = relative_path_in_destination_storage;
-    processed_part_entry.status = "SUCCESS";
     processed_part_entry.finished_by = storage.replica_name;
 
     requests.emplace_back(zkutil::makeRemoveRequest(processing_parts_path / part_name, -1));

src/Storages/MergeTree/ExportPartitionUtils.cpp

@@ -2,6 +2,7 @@
 #include <Common/ZooKeeper/ZooKeeper.h>
 #include <Common/logger_useful.h>
 #include "Storages/ExportReplicatedMergeTreePartitionManifest.h"
+#include "Storages/ExportReplicatedMergeTreePartitionTaskEntry.h"
 #include <filesystem>
 
 namespace DB
@@ -80,7 +81,7 @@ namespace ExportPartitionUtils
         destination_storage->commitExportPartitionTransaction(manifest.transaction_id, manifest.partition_id, exported_paths, context);
 
         LOG_INFO(log, "ExportPartition: Committed export, mark as completed");
-        if (Coordination::Error::ZOK == zk->trySet(fs::path(entry_path) / "status", "COMPLETED", -1))
+        if (Coordination::Error::ZOK == zk->trySet(fs::path(entry_path) / "status", String(magic_enum::enum_name(ExportReplicatedMergeTreePartitionTaskEntry::Status::COMPLETED)).data(), -1))
         {
             LOG_INFO(log, "ExportPartition: Marked export as completed");
         }

src/Storages/MergeTree/tests/gtest_export_partition_ordering.cpp

@@ -40,9 +40,9 @@ TEST_F(ExportPartitionOrderingTest, IterationOrderMatchesCreateTime)
     manifest3.destination_table = "table1";
     manifest3.create_time = base_time; // Oldest
 
-    ExportReplicatedMergeTreePartitionTaskEntry entry1{manifest1, {}};
-    ExportReplicatedMergeTreePartitionTaskEntry entry2{manifest2, {}};
-    ExportReplicatedMergeTreePartitionTaskEntry entry3{manifest3, {}};
+    ExportReplicatedMergeTreePartitionTaskEntry entry1{manifest1, ExportReplicatedMergeTreePartitionTaskEntry::Status::PENDING, {}};
+    ExportReplicatedMergeTreePartitionTaskEntry entry2{manifest2, ExportReplicatedMergeTreePartitionTaskEntry::Status::PENDING, {}};
+    ExportReplicatedMergeTreePartitionTaskEntry entry3{manifest3, ExportReplicatedMergeTreePartitionTaskEntry::Status::PENDING, {}};
 
     // Insert in reverse order
     by_key.insert(entry1);

src/Storages/StorageReplicatedMergeTree.cpp

@@ -8191,7 +8191,7 @@ void StorageReplicatedMergeTree::exportPartitionToTable(const PartitionCommand &
     for (const auto & part : part_names)
     {
         ExportReplicatedMergeTreePartitionProcessingPartEntry entry;
-        entry.status = "PENDING";
+        entry.status = ExportReplicatedMergeTreePartitionProcessingPartEntry::Status::PENDING;
         entry.part_name = part;
         entry.retry_count = 0;
 

tests/integration/test_export_replicated_mt_partition_to_object_storage/test.py

@@ -553,16 +553,15 @@ def test_export_ttl(cluster):
     time.sleep(expiration_time)
 
     # assert that the export succeeded, check the commit file
-    assert node.query(f"SELECT count() FROM s3(s3_conn, filename='{s3_table}/commit_2020_*', format=LineAsString)") == '1\n', "Export did not succeed"
+    # there should be two commit files now, one for the first export and one for the second export
+    assert node.query(f"SELECT count() FROM s3(s3_conn, filename='{s3_table}/commit_2020_*', format=LineAsString)") == '2\n', "Export did not succeed"
 
 
-# export an individual part with alter table export part
-# and then try to export the partition. It should not fail because export partition is idempotent.
-def test_export_part_and_partition(cluster):
+def test_export_partition_file_already_exists_policy(cluster):
     node = cluster.instances["replica1"]
 
-    mt_table = "export_part_and_partition_mt_table"
-    s3_table = "export_part_and_partition_s3_table"
+    mt_table = "export_partition_file_already_exists_policy_mt_table"
+    s3_table = "export_partition_file_already_exists_policy_s3_table"
 
     create_tables_and_insert_data(node, mt_table, s3_table, "replica1")