Yet another export replicated partition pr

src/Access/Common/AccessType.h

@@ -211,6 +211,7 @@ enum class AccessType : uint8_t
     M(ALTER_SETTINGS, "ALTER SETTING, ALTER MODIFY SETTING, MODIFY SETTING, RESET SETTING", TABLE, ALTER_TABLE) /* allows to execute ALTER MODIFY SETTING */\
     M(ALTER_MOVE_PARTITION, "ALTER MOVE PART, MOVE PARTITION, MOVE PART", TABLE, ALTER_TABLE) \
     M(ALTER_EXPORT_PART, "ALTER EXPORT PART, EXPORT PART", TABLE, ALTER_TABLE) \
+    M(ALTER_EXPORT_PARTITION, "ALTER EXPORT PARTITION, EXPORT PARTITION", TABLE, ALTER_TABLE) \
     M(ALTER_FETCH_PARTITION, "ALTER FETCH PART, FETCH PARTITION", TABLE, ALTER_TABLE) \
     M(ALTER_FREEZE_PARTITION, "FREEZE PARTITION, UNFREEZE", TABLE, ALTER_TABLE) \
     M(ALTER_UNLOCK_SNAPSHOT, "UNLOCK SNAPSHOT", TABLE, ALTER_TABLE) \

src/Core/ServerSettings.cpp

@@ -1144,6 +1144,7 @@ The policy on how to perform a scheduling of CPU slots specified by `concurrent_
     DECLARE(UInt64, object_storage_list_objects_cache_max_entries, 1000, "Maximum size of ObjectStorage list objects cache in entries. Zero means disabled.", 0)                      \
     DECLARE(UInt64, object_storage_list_objects_cache_ttl, 3600, "Time to live of records in ObjectStorage list objects cache in seconds. Zero means unlimited", 0)                   \
     DECLARE(UInt64, input_format_parquet_metadata_cache_max_size, 500000000, "Maximum size of parquet file metadata cache", 0) \
+    DECLARE(Bool, enable_experimental_export_merge_tree_partition_feature, false, "Enable export replicated merge tree partition feature. It is experimental and not yet ready for production use.", 0) \
 // clang-format on
 
 /// If you add a setting which can be updated at runtime, please update 'changeable_settings' map in dumpToSystemServerSettingsColumns below

src/Core/Settings.cpp

@@ -6873,6 +6873,22 @@ Use roaring bitmap for iceberg positional deletes.
 )", 0) \
     DECLARE(Bool, export_merge_tree_part_overwrite_file_if_exists, false, R"(
 Overwrite file if it already exists when exporting a merge tree part
+)", 0) \
+    DECLARE(Bool, export_merge_tree_partition_force_export, false, R"(
+Ignore existing partition export and overwrite the zookeeper entry
+)", 0) \
+    DECLARE(UInt64, export_merge_tree_partition_max_retries, 3, R"(
+Maximum number of retries for exporting a merge tree part in an export partition task
+)", 0) \
+    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(MergeTreePartExportFileAlreadyExistsPolicy, export_merge_tree_part_file_already_exists_policy, MergeTreePartExportFileAlreadyExistsPolicy::skip, R"(
+Possible values:
+- skip - Skip the file if it already exists.
+- error - Throw an error if the file already exists.
+- overwrite - Overwrite the file.
 )", 0) \
     DECLARE(Timezone, iceberg_timezone_for_timestamptz, "UTC", R"(
 Timezone for Iceberg timestamptz field.

src/Core/Settings.h

@@ -81,6 +81,7 @@ class WriteBuffer;
     M(CLASS_NAME, LogsLevel) \
     M(CLASS_NAME, Map) \
     M(CLASS_NAME, MaxThreads) \
+    M(CLASS_NAME, MergeTreePartExportFileAlreadyExistsPolicy) \
     M(CLASS_NAME, Milliseconds) \
     M(CLASS_NAME, MsgPackUUIDRepresentation) \
     M(CLASS_NAME, MySQLDataTypesSupport) \

src/Core/SettingsChangesHistory.cpp

@@ -47,6 +47,10 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
             {"allow_retries_in_cluster_requests", false, false, "New setting"},
             {"object_storage_remote_initiator", false, false, "New setting."},
             {"allow_experimental_export_merge_tree_part", false, true, "Turned ON by default for Antalya."},
+            {"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", "skip", "skip", "New setting."},
             {"iceberg_timezone_for_timestamptz", "UTC", "UTC", "New setting."}
         });
         addSettingsChanges(settings_changes_history, "25.8",

src/Core/SettingsEnums.cpp

@@ -370,4 +370,6 @@ IMPLEMENT_SETTING_ENUM(
      {"manifest_list_entry", IcebergMetadataLogLevel::ManifestListEntry},
      {"manifest_file_metadata", IcebergMetadataLogLevel::ManifestFileMetadata},
      {"manifest_file_entry", IcebergMetadataLogLevel::ManifestFileEntry}})
+
+IMPLEMENT_SETTING_AUTO_ENUM(MergeTreePartExportFileAlreadyExistsPolicy, ErrorCodes::BAD_ARGUMENTS);
 }

src/Core/SettingsEnums.h

@@ -480,4 +480,14 @@ enum class IcebergMetadataLogLevel : uint8_t
 };
 
 DECLARE_SETTING_ENUM(IcebergMetadataLogLevel)
+
+enum class MergeTreePartExportFileAlreadyExistsPolicy : uint8_t
+{
+    skip,
+    error,
+    overwrite,
+};
+
+DECLARE_SETTING_ENUM(MergeTreePartExportFileAlreadyExistsPolicy)
+
 }

src/Functions/generateSnowflakeID.cpp

@@ -154,6 +154,11 @@ uint64_t generateSnowflakeID()
     return fromSnowflakeId(snowflake_id);
 }
 
+std::string generateSnowflakeIDString()
+{
+    return std::to_string(generateSnowflakeID());
+}
+
 class FunctionGenerateSnowflakeID : public IFunction
 {
 public:

src/Functions/generateSnowflakeID.h

@@ -7,4 +7,6 @@ namespace DB
 
 uint64_t generateSnowflakeID();
 
+std::string generateSnowflakeIDString();
+
 }

src/Interpreters/InterpreterAlterQuery.cpp

@@ -545,6 +545,12 @@ AccessRightsElements InterpreterAlterQuery::getRequiredAccessForCommand(const AS
             required_access.emplace_back(AccessType::INSERT, command.to_database, command.to_table);
             break;
         }
+        case ASTAlterCommand::EXPORT_PARTITION:
+        {
+            required_access.emplace_back(AccessType::ALTER_EXPORT_PARTITION, database, table);
+            required_access.emplace_back(AccessType::INSERT, command.to_database, command.to_table);
+            break;
+        }
         case ASTAlterCommand::FETCH_PARTITION:
         {
             required_access.emplace_back(AccessType::ALTER_FETCH_PARTITION, database, table);

src/Interpreters/InterpreterKillQueryQuery.cpp

@@ -19,6 +19,7 @@
 #include <DataTypes/DataTypeString.h>
 #include <Columns/ColumnsNumber.h>
 #include <DataTypes/DataTypesNumber.h>
+#include <Core/ServerSettings.h>
 #include <Processors/Sources/SourceFromSingleChunk.h>
 #include <Processors/ISource.h>
 #include <Processors/Executors/PullingPipelineExecutor.h>
@@ -37,11 +38,17 @@ namespace Setting
     extern const SettingsUInt64 max_parser_depth;
 }
 
+namespace ServerSetting
+{
+    extern const ServerSettingsBool enable_experimental_export_merge_tree_partition_feature;
+}
+
 namespace ErrorCodes
 {
     extern const int LOGICAL_ERROR;
     extern const int ACCESS_DENIED;
     extern const int NOT_IMPLEMENTED;
+    extern const int SUPPORT_IS_DISABLED;
 }
 
 
@@ -250,6 +257,82 @@ BlockIO InterpreterKillQueryQuery::execute()
 
         break;
     }
+    case ASTKillQueryQuery::Type::ExportPartition:
+    {
+        if (!getContext()->getServerSettings()[ServerSetting::enable_experimental_export_merge_tree_partition_feature])
+        {
+            throw Exception(ErrorCodes::SUPPORT_IS_DISABLED,
+                "Exporting merge tree partition is experimental. Set the server setting `enable_experimental_export_merge_tree_partition_feature` to enable it");
+        }
+
+        Block exports_block = getSelectResult(
+            "source_database, source_table, transaction_id, destination_database, destination_table, partition_id",
+            "system.replicated_partition_exports");
+        if (exports_block.empty())
+            return res_io;
+
+        const ColumnString & src_db_col = typeid_cast<const ColumnString &>(*exports_block.getByName("source_database").column);
+        const ColumnString & src_table_col = typeid_cast<const ColumnString &>(*exports_block.getByName("source_table").column);
+        const ColumnString & dst_db_col = typeid_cast<const ColumnString &>(*exports_block.getByName("destination_database").column);
+        const ColumnString & dst_table_col = typeid_cast<const ColumnString &>(*exports_block.getByName("destination_table").column);
+        const ColumnString & tx_col = typeid_cast<const ColumnString &>(*exports_block.getByName("transaction_id").column);
+
+        auto header = exports_block.cloneEmpty();
+        header.insert(0, {ColumnString::create(), std::make_shared<DataTypeString>(), "kill_status"});
+
+        MutableColumns res_columns = header.cloneEmptyColumns();
+        AccessRightsElements required_access_rights;
+        auto access = getContext()->getAccess();
+        bool access_denied = false;
+
+        for (size_t i = 0; i < exports_block.rows(); ++i)
+        {
+            const auto src_database = src_db_col.getDataAt(i).toString();
+            const auto src_table = src_table_col.getDataAt(i).toString();
+            const auto dst_database = dst_db_col.getDataAt(i).toView();
+            const auto dst_table = dst_table_col.getDataAt(i).toView();
+
+            const auto table_id = StorageID{src_database, src_table};
+            const auto transaction_id = tx_col.getDataAt(i).toString();
+
+            CancellationCode code = CancellationCode::Unknown;
+            if (!query.test)
+            {
+                auto storage = DatabaseCatalog::instance().tryGetTable(table_id, getContext());
+                if (!storage)
+                    code = CancellationCode::NotFound;
+                else
+                {
+                    ASTAlterCommand alter_command{};
+                    alter_command.type = ASTAlterCommand::EXPORT_PARTITION;
+                    alter_command.move_destination_type = DataDestinationType::TABLE;
+                    alter_command.from_database = src_database;
+                    alter_command.from_table = src_table;
+                    alter_command.to_database = dst_database;
+                    alter_command.to_table = dst_table;
+
+                    required_access_rights = InterpreterAlterQuery::getRequiredAccessForCommand(
+                        alter_command, table_id.database_name, table_id.table_name);
+                    if (!access->isGranted(required_access_rights))
+                    {
+                        access_denied = true;
+                        continue;
+                    }
+                    code = storage->killExportPartition(transaction_id);
+                }
+            }
+
+            insertResultRow(i, code, exports_block, header, res_columns);
+        }
+
+        if (res_columns[0]->empty() && access_denied)
+            throw Exception(ErrorCodes::ACCESS_DENIED, "Not allowed to kill export partition. "
+                "To execute this query, it's necessary to have the grant {}", required_access_rights.toString());
+
+        res_io.pipeline = QueryPipeline(Pipe(std::make_shared<SourceFromSingleChunk>(std::make_shared<const Block>(header.cloneWithColumns(std::move(res_columns))))));
+
+        break;
+    }
     case ASTKillQueryQuery::Type::Mutation:
     {
         Block mutations_block = getSelectResult("database, table, mutation_id, command", "system.mutations");
@@ -462,6 +545,9 @@ AccessRightsElements InterpreterKillQueryQuery::getRequiredAccessForDDLOnCluster
                 | AccessType::ALTER_MATERIALIZE_COLUMN
                 | AccessType::ALTER_MATERIALIZE_TTL
             );
+    /// todo arthur think about this
+    else if (query.type == ASTKillQueryQuery::Type::ExportPartition)
+        required_access.emplace_back(AccessType::ALTER_EXPORT_PARTITION);
     return required_access;
 }
 

src/Parsers/ASTAlterQuery.cpp

@@ -378,6 +378,17 @@ void ASTAlterCommand::formatImpl(WriteBuffer & ostr, const FormatSettings & sett
         }
 
     }
+    else if (type == ASTAlterCommand::EXPORT_PARTITION)
+    {
+        ostr << "EXPORT PARTITION ";
+        partition->format(ostr, settings, state, frame);
+        ostr << " TO TABLE ";
+        if (!to_database.empty())
+        {
+            ostr << backQuoteIfNeed(to_database) << ".";
+        }
+        ostr << backQuoteIfNeed(to_table);
+    }
     else if (type == ASTAlterCommand::REPLACE_PARTITION)
     {
         ostr << (replace ? "REPLACE" : "ATTACH") << " PARTITION "

src/Parsers/ASTAlterQuery.h

@@ -72,6 +72,7 @@ class ASTAlterCommand : public IAST
         UNFREEZE_PARTITION,
         UNFREEZE_ALL,
         EXPORT_PART,
+        EXPORT_PARTITION,
 
         DELETE,
         UPDATE,

src/Parsers/ASTKillQueryQuery.cpp

@@ -27,6 +27,9 @@ void ASTKillQueryQuery::formatQueryImpl(WriteBuffer & ostr, const FormatSettings
         case Type::Transaction:
             ostr << "TRANSACTION";
             break;
+        case Type::ExportPartition:
+            ostr << "EXPORT PARTITION";
+            break;
     }
 
     formatOnCluster(ostr, settings);

src/Parsers/ASTKillQueryQuery.h

@@ -13,6 +13,7 @@ class ASTKillQueryQuery : public ASTQueryWithOutput, public ASTQueryWithOnCluste
     {
         Query,      /// KILL QUERY
         Mutation,   /// KILL MUTATION
+        ExportPartition, /// KILL EXPORT_PARTITION
         PartMoveToShard, /// KILL PART_MOVE_TO_SHARD
         Transaction,     /// KILL TRANSACTION
     };

src/Parsers/CommonParsers.h

@@ -333,6 +333,7 @@ namespace DB
     MR_MACROS(MOVE_PART, "MOVE PART") \
     MR_MACROS(MOVE_PARTITION, "MOVE PARTITION") \
     MR_MACROS(EXPORT_PART, "EXPORT PART") \
+    MR_MACROS(EXPORT_PARTITION, "EXPORT PARTITION") \
     MR_MACROS(MOVE, "MOVE") \
     MR_MACROS(MS, "MS") \
     MR_MACROS(MUTATION, "MUTATION") \

src/Parsers/ParserAlterQuery.cpp

@@ -83,6 +83,7 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected
     ParserKeyword s_move_partition(Keyword::MOVE_PARTITION);
     ParserKeyword s_move_part(Keyword::MOVE_PART);
     ParserKeyword s_export_part(Keyword::EXPORT_PART);
+    ParserKeyword s_export_partition(Keyword::EXPORT_PARTITION);
     ParserKeyword s_drop_detached_partition(Keyword::DROP_DETACHED_PARTITION);
     ParserKeyword s_drop_detached_part(Keyword::DROP_DETACHED_PART);
     ParserKeyword s_fetch_partition(Keyword::FETCH_PARTITION);
@@ -553,6 +554,22 @@ bool ParserAlterCommand::parseImpl(Pos & pos, ASTPtr & node, Expected & expected
                     return false;
                 command->move_destination_type = DataDestinationType::TABLE;
             }
+            else if (s_export_partition.ignore(pos, expected))
+            {
+                if (!parser_partition.parse(pos, command_partition, expected))
+                    return false;
+
+                command->type = ASTAlterCommand::EXPORT_PARTITION;
+
+                if (!s_to_table.ignore(pos, expected))
+                {
+                    return false;
+                }
+
+                if (!parseDatabaseAndTableName(pos, expected, command->to_database, command->to_table))
+                    return false;
+                command->move_destination_type = DataDestinationType::TABLE;
+            }
             else if (s_move_partition.ignore(pos, expected))
             {
                 if (!parser_partition.parse(pos, command_partition, expected))

src/Parsers/ParserKillQueryQuery.cpp

@@ -17,6 +17,7 @@ bool ParserKillQueryQuery::parseImpl(Pos & pos, ASTPtr & node, Expected & expect
     ParserKeyword p_kill{Keyword::KILL};
     ParserKeyword p_query{Keyword::QUERY};
     ParserKeyword p_mutation{Keyword::MUTATION};
+    ParserKeyword p_export_partition{Keyword::EXPORT_PARTITION};
     ParserKeyword p_part_move_to_shard{Keyword::PART_MOVE_TO_SHARD};
     ParserKeyword p_transaction{Keyword::TRANSACTION};
     ParserKeyword p_on{Keyword::ON};
@@ -33,6 +34,8 @@ bool ParserKillQueryQuery::parseImpl(Pos & pos, ASTPtr & node, Expected & expect
         query->type = ASTKillQueryQuery::Type::Query;
     else if (p_mutation.ignore(pos, expected))
         query->type = ASTKillQueryQuery::Type::Mutation;
+    else if (p_export_partition.ignore(pos, expected))
+        query->type = ASTKillQueryQuery::Type::ExportPartition;
     else if (p_part_move_to_shard.ignore(pos, expected))
         query->type = ASTKillQueryQuery::Type::PartMoveToShard;
     else if (p_transaction.ignore(pos, expected))