diff --git a/src/Interpreters/InterpreterKillQueryQuery.cpp b/src/Interpreters/InterpreterKillQueryQuery.cpp index 4ab33f67cb75..0dfe78127848 100644 --- a/src/Interpreters/InterpreterKillQueryQuery.cpp +++ b/src/Interpreters/InterpreterKillQueryQuery.cpp @@ -265,19 +265,35 @@ BlockIO InterpreterKillQueryQuery::execute() "Exporting merge tree partition is experimental. Set the server setting `allow_experimental_export_merge_tree_partition` 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 String export_columns = "source_database, source_table, transaction_id, destination_database, destination_table, partition_id"; + + /// Partition exports live in two system tables: `replicated_partition_exports` for + /// ReplicatedMergeTree and `partition_exports` for plain MergeTree. Query both so a single + /// KILL EXPORT PARTITION targets either engine. Each query applies the user WHERE to its own + /// table (which exposes all of its columns); a WHERE that references columns specific to the + /// other engine simply yields no matches / is tolerated. + Block replicated_exports_block = getSelectResult(export_columns, "system.replicated_partition_exports"); + Block plain_exports_block; + try + { + plain_exports_block = getSelectResult(export_columns, "system.partition_exports"); + } + catch (...) + { + tryLogCurrentException(getLogger("InterpreterKillQueryQuery"), + "KILL EXPORT PARTITION: could not read system.partition_exports (the WHERE may reference " + "columns that only exist for ReplicatedMergeTree); ignoring plain MergeTree tables"); + } - const ColumnString & src_db_col = typeid_cast(*exports_block.getByName("source_database").column); - const ColumnString & src_table_col = typeid_cast(*exports_block.getByName("source_table").column); - const ColumnString & dst_db_col = typeid_cast(*exports_block.getByName("destination_database").column); - const ColumnString & dst_table_col = typeid_cast(*exports_block.getByName("destination_table").column); - const ColumnString & tx_col = typeid_cast(*exports_block.getByName("transaction_id").column); + if (replicated_exports_block.empty() && plain_exports_block.empty()) + return res_io; - auto header = exports_block.cloneEmpty(); + /// Build the result header explicitly from the fixed projection so it does not depend on + /// whether either source block ended up with rows. + Block header; + for (const auto & column_name : {"source_database", "source_table", "transaction_id", + "destination_database", "destination_table", "partition_id"}) + header.insert({ColumnString::create(), std::make_shared(), column_name}); header.insert(0, {ColumnString::create(), std::make_shared(), "kill_status"}); MutableColumns res_columns = header.cloneEmptyColumns(); @@ -285,45 +301,60 @@ BlockIO InterpreterKillQueryQuery::execute() auto access = getContext()->getAccess(); bool access_denied = false; - for (size_t i = 0; i < exports_block.rows(); ++i) + auto process_block = [&](const Block & exports_block) { - const auto src_database = src_db_col.getDataAt(i); - const auto src_table = src_table_col.getDataAt(i); - const auto dst_database = dst_db_col.getDataAt(i); - const auto dst_table = dst_table_col.getDataAt(i); + if (exports_block.empty()) + return; - const auto table_id = StorageID{std::string{src_database}, std::string{src_table}}; - const auto transaction_id = tx_col.getDataAt(i); + const ColumnString & src_db_col = typeid_cast(*exports_block.getByName("source_database").column); + const ColumnString & src_table_col = typeid_cast(*exports_block.getByName("source_table").column); + const ColumnString & dst_db_col = typeid_cast(*exports_block.getByName("destination_database").column); + const ColumnString & dst_table_col = typeid_cast(*exports_block.getByName("destination_table").column); + const ColumnString & tx_col = typeid_cast(*exports_block.getByName("transaction_id").column); - CancellationCode code = CancellationCode::Unknown; - if (!query.test) + for (size_t i = 0; i < exports_block.rows(); ++i) { - 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; + const auto src_database = src_db_col.getDataAt(i); + const auto src_table = src_table_col.getDataAt(i); + const auto dst_database = dst_db_col.getDataAt(i); + const auto dst_table = dst_table_col.getDataAt(i); - required_access_rights = InterpreterAlterQuery::getRequiredAccessForCommand( - alter_command, table_id.database_name, table_id.table_name); - if (!access->isGranted(required_access_rights)) + const auto table_id = StorageID{std::string{src_database}, std::string{src_table}}; + const auto transaction_id = tx_col.getDataAt(i); + + CancellationCode code = CancellationCode::Unknown; + if (!query.test) + { + auto storage = DatabaseCatalog::instance().tryGetTable(table_id, getContext()); + if (!storage) + code = CancellationCode::NotFound; + else { - access_denied = true; - continue; + 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(std::string{transaction_id}); } - code = storage->killExportPartition(std::string{transaction_id}); } + + insertResultRow(i, code, exports_block, header, res_columns); } + }; - insertResultRow(i, code, exports_block, header, res_columns); - } + process_block(replicated_exports_block); + process_block(plain_exports_block); if (res_columns[0]->empty() && access_denied) throw Exception(ErrorCodes::ACCESS_DENIED, "Not allowed to kill export partition. " diff --git a/src/Storages/MergeTree/ExportPartitionUtils.cpp b/src/Storages/MergeTree/ExportPartitionUtils.cpp index ec8046f3e991..9777222d36a1 100644 --- a/src/Storages/MergeTree/ExportPartitionUtils.cpp +++ b/src/Storages/MergeTree/ExportPartitionUtils.cpp @@ -7,6 +7,10 @@ #include "Storages/ExportReplicatedMergeTreePartitionManifest.h" #include "Storages/ExportReplicatedMergeTreePartitionTaskEntry.h" #include +#include +#include +#include +#include #include #include #include @@ -20,6 +24,7 @@ #if USE_AVRO #include #include +#include #endif namespace ProfileEvents @@ -75,6 +80,7 @@ namespace ErrorCodes namespace Setting { extern const SettingsBool export_merge_tree_part_allow_lossy_cast; + extern const SettingsBool allow_insert_into_iceberg; } namespace FailPoints @@ -147,7 +153,8 @@ namespace ExportPartitionUtils return parts.front()->minmax_idx->getBlock(storage); } - ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ExportReplicatedMergeTreePartitionManifest & manifest) + template + ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ManifestT & manifest) { auto context_copy = Context::createCopy(context); context_copy->makeQueryContextForExportPart(); @@ -187,6 +194,102 @@ namespace ExportPartitionUtils return context_copy; } + template ContextPtr getContextCopyWithTaskSettings( + const ContextPtr &, const ExportReplicatedMergeTreePartitionManifest &); + template ContextPtr getContextCopyWithTaskSettings( + const ContextPtr &, const MergeTreePartitionExportTask &); + + std::string extractDestinationIcebergMetadataJson( + const StorageMetadataPtr & source_metadata, + const StorageID & source_storage_id, + const StoragePtr & dest_storage, + const ContextPtr & context) + { + if (dest_storage->getStorageID() == source_storage_id) + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Exporting to the same table is not allowed"); + + if (!dest_storage->supportsImport(context)) + throw Exception(ErrorCodes::NOT_IMPLEMENTED, + "Destination storage {} does not support MergeTree parts or uses unsupported partitioning", + dest_storage->getName()); + + const auto destination_metadata = dest_storage->getInMemoryMetadataPtr(); + + /// Positional CAST matching, like `INSERT INTO dest SELECT * FROM src`. + verifyExportSchemaCastable(source_metadata, destination_metadata, dest_storage->getStorageID(), context); + + auto query_to_string = [](const ASTPtr & ast) { return ast ? ast->formatWithSecretsOneLine() : ""; }; + + /// Iceberg partition compatibility is checked below; for non-data-lake targets we only need + /// the partition-key ASTs to match (partition-column types follow the lossy-cast gate). + if (!dest_storage->isDataLake()) + { + if (query_to_string(source_metadata->getPartitionKeyAST()) != query_to_string(destination_metadata->getPartitionKeyAST())) + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Tables have different partition key"); + return {}; + } + +#if USE_AVRO + auto * object_storage = dynamic_cast(dest_storage.get()); + auto * object_storage_cluster = dynamic_cast(dest_storage.get()); + + /// in theory this should never happen, but just in case + if (!object_storage && !object_storage_cluster) + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Destination storage {} is not a StorageObjectStorage", dest_storage->getName()); + + IcebergMetadata * iceberg_metadata = nullptr; + if (object_storage) + iceberg_metadata = dynamic_cast(object_storage->getExternalMetadata(context)); + else if (object_storage_cluster) + iceberg_metadata = dynamic_cast(object_storage_cluster->getExternalMetadata(context)); + if (!iceberg_metadata) + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Destination storage {} is a data lake but not an iceberg table", dest_storage->getName()); + + if (!context->getSettingsRef()[Setting::allow_insert_into_iceberg]) + throw Exception(ErrorCodes::SUPPORT_IS_DISABLED, + "Iceberg writes are experimental. " + "To allow its usage, enable the setting `allow_insert_into_iceberg`."); + + const auto metadata_object = iceberg_metadata->getMetadataJSON(context); + + verifyIcebergPartitionCompatibility(metadata_object, source_metadata->getPartitionKeyAST()); + + std::ostringstream oss; // STYLE_CHECK_ALLOW_STD_STRING_STREAM + oss.exceptions(std::ios::failbit); + metadata_object->stringify(oss); + return oss.str(); +#else + throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Data lake export requires Avro support"); +#endif + } + + void commitExportedPaths( + const String & transaction_id, + const String & partition_id, + const String & iceberg_metadata_json, + bool write_full_path_in_iceberg_metadata, + const std::vector & exported_paths, + const StoragePtr & destination_storage, + MergeTreeData & source_storage, + const ContextPtr & context_in) + { + auto context = Context::createCopy(context_in); + context->setSetting("write_full_path_in_iceberg_metadata", write_full_path_in_iceberg_metadata); + + IStorage::IcebergCommitExportPartitionArguments iceberg_args; + + if (!iceberg_metadata_json.empty()) + { + iceberg_args.metadata_json_string = iceberg_metadata_json; + if (source_storage.getInMemoryMetadataPtr()->hasPartitionKey()) + iceberg_args.partition_source_block = + getPartitionSourceBlockForIcebergCommit(source_storage, partition_id); + } + + destination_storage->commitExportPartitionTransaction( + transaction_id, partition_id, exported_paths, iceberg_args, context); + } + /// Collect all the exported paths from the processed parts /// If multiRead is supported by the keeper implementation, it is done in a single request /// Otherwise, multiple async requests are sent @@ -253,9 +356,6 @@ namespace ExportPartitionUtils MergeTreeData & source_storage, const String & replica_name) { - auto context = Context::createCopy(context_in); - context->setSetting("write_full_path_in_iceberg_metadata", manifest.write_full_path_in_iceberg_metadata); - /// Failpoint used by integration tests to force persistent commit failure and exercise /// the commit-attempts budget / FAILED state transition. fiu_do_on(FailPoints::export_partition_commit_always_throw, @@ -302,17 +402,15 @@ namespace ExportPartitionUtils throw Exception(ErrorCodes::CORRUPTED_DATA, "ExportPartition: Reached the commit phase, but exported paths size is less than the number of parts, will not commit export. This might be a bug"); } - IStorage::IcebergCommitExportPartitionArguments iceberg_args; - - if (!manifest.iceberg_metadata_json.empty()) - { - iceberg_args.metadata_json_string = manifest.iceberg_metadata_json; - if (source_storage.getInMemoryMetadataPtr()->hasPartitionKey()) - iceberg_args.partition_source_block = - getPartitionSourceBlockForIcebergCommit(source_storage, manifest.partition_id); - } - - destination_storage->commitExportPartitionTransaction(manifest.transaction_id, manifest.partition_id, exported_paths, iceberg_args, context); + commitExportedPaths( + manifest.transaction_id, + manifest.partition_id, + manifest.iceberg_metadata_json, + manifest.write_full_path_in_iceberg_metadata, + exported_paths, + destination_storage, + source_storage, + context_in); /// Failpoint to simulate a crash after the Iceberg commit succeeds but before /// ZooKeeper is updated to COMPLETED. Used by idempotency integration tests. diff --git a/src/Storages/MergeTree/ExportPartitionUtils.h b/src/Storages/MergeTree/ExportPartitionUtils.h index 0bb8acb9bda4..3696226feb78 100644 --- a/src/Storages/MergeTree/ExportPartitionUtils.h +++ b/src/Storages/MergeTree/ExportPartitionUtils.h @@ -27,7 +27,36 @@ namespace ExportPartitionUtils std::vector getExportedPaths(const LoggerPtr & log, const zkutil::ZooKeeperPtr & zk, const std::string & export_path); - ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ExportReplicatedMergeTreePartitionManifest & manifest); + /// Build a query context carrying the export task's persisted settings. Templated on the + /// descriptor type so it serves both the replicated manifest (backed by ZooKeeper) and the + /// plain `MergeTreePartitionExportTask` (backed by disk); both expose the same setting fields. + template + ContextPtr getContextCopyWithTaskSettings(const ContextPtr & context, const ManifestT & manifest); + + /// Validates that `dest_storage` is a legal export target for a source table with the given + /// metadata: it must support imports, be positionally castable, and (for non-data-lake targets) + /// share the same partition key. For Iceberg destinations it additionally verifies partition-spec + /// compatibility and returns the serialized destination `metadata.json`; returns an empty string + /// for non-data-lake destinations. Throws on any incompatibility. + std::string extractDestinationIcebergMetadataJson( + const StorageMetadataPtr & source_metadata, + const StorageID & source_storage_id, + const StoragePtr & dest_storage, + const ContextPtr & context); + + /// ZooKeeper-free commit core shared by the replicated and plain partition-export paths: + /// assembles the Iceberg commit arguments (deriving the partition source block from the source + /// table when the destination is a data lake with a partition key) and invokes + /// `commitExportPartitionTransaction` on the destination storage. + void commitExportedPaths( + const String & transaction_id, + const String & partition_id, + const String & iceberg_metadata_json, + bool write_full_path_in_iceberg_metadata, + const std::vector & exported_paths, + const StoragePtr & destination_storage, + MergeTreeData & source_storage, + const ContextPtr & context); /// Returns the representative source partition-key columns (the first active local part's /// minmax block) for the given partition_id. The destination recomputes the Iceberg partition diff --git a/src/Storages/MergeTree/MergeTreePartitionExportScheduler.cpp b/src/Storages/MergeTree/MergeTreePartitionExportScheduler.cpp new file mode 100644 index 000000000000..18a9ed3a7133 --- /dev/null +++ b/src/Storages/MergeTree/MergeTreePartitionExportScheduler.cpp @@ -0,0 +1,569 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace fs = std::filesystem; + +namespace DB +{ + +namespace ErrorCodes +{ + extern const int EXPORT_PARTITION_ALREADY_EXPORTED; + extern const int CORRUPTED_DATA; + extern const int QUERY_WAS_CANCELLED; + extern const int UNKNOWN_TABLE; +} + +MergeTreePartitionExportScheduler::MergeTreePartitionExportScheduler(StorageMergeTree & storage_) + : storage(storage_) +{ +} + +String MergeTreePartitionExportScheduler::compositeKey( + const String & partition_id, const String & destination_database, const String & destination_table) +{ + const QualifiedTableName qualified_table_name{destination_database, destination_table}; + return partition_id + "_" + qualified_table_name.getFullName(); +} + +String MergeTreePartitionExportScheduler::getExportsRelativePath() const +{ + return fs::path(storage.getRelativeDataPath()) / "partition_exports"; +} + +void MergeTreePartitionExportScheduler::throwIfAlreadyExported(const String & composite_key, bool force) const +{ + if (force) + return; + + std::lock_guard lock(mutex); + if (tasks.contains(composite_key)) + throw Exception(ErrorCodes::EXPORT_PARTITION_ALREADY_EXPORTED, + "Export with key {} already exported or it is being exported. " + "Set `export_merge_tree_partition_force_export` to overwrite it.", + composite_key); +} + +void MergeTreePartitionExportScheduler::addTask( + MergeTreePartitionExportTask descriptor, std::vector part_references, bool force) +{ + const auto composite_key = compositeKey(descriptor.partition_id, descriptor.destination_database, descriptor.destination_table); + const auto transaction_id = descriptor.transaction_id; + String descriptor_json; + + String previous_transaction_id; + bool replaced = false; + + { + std::lock_guard lock(mutex); + + if (const auto it = tasks.find(composite_key); it != tasks.end()) + { + if (!force) + throw Exception(ErrorCodes::EXPORT_PARTITION_ALREADY_EXPORTED, + "Export with key {} already exported or it is being exported. " + "Set `export_merge_tree_partition_force_export` to overwrite it.", + composite_key); + + /// Overwrite: erase the previous task atomically here; its in-flight part exports and + /// its on-disk file are cleaned up below, outside the registry lock (see note). + previous_transaction_id = it->second.descriptor.transaction_id; + replaced = true; + tasks.erase(it); + } + + TaskEntry entry; + entry.descriptor = std::move(descriptor); + entry.part_references = std::move(part_references); + descriptor_json = entry.descriptor.toJsonString(); + tasks.emplace(composite_key, std::move(entry)); + } + + /// killExportPart takes export_manifests_mutex, and a part-export completion callback runs while + /// holding export_manifests_mutex and then takes our registry mutex. Calling killExportPart while + /// holding the registry mutex would therefore be an AB-BA lock inversion, so we do it here after + /// releasing the lock. The previous task's entry is already gone, so its late completion callbacks + /// find nothing and no-op. + if (replaced) + { + LOG_INFO(storage.log, "ExportPartition: overwriting export with key {}", composite_key); + storage.killExportPart(previous_transaction_id); + removeTaskFile(previous_transaction_id); + } + + /// Persist before announcing success so a crash right after this call leaves a resumable task. + persist(transaction_id, descriptor_json); + + LOG_INFO(storage.log, "ExportPartition: scheduled export task {} (key {})", transaction_id, composite_key); + storage.triggerPartitionExportTask(); +} + +CancellationCode MergeTreePartitionExportScheduler::kill(const String & transaction_id) +{ + String descriptor_json; + { + std::lock_guard lock(mutex); + TaskEntry * target = nullptr; + for (auto & [key, entry] : tasks) + { + if (entry.descriptor.transaction_id == transaction_id) + { + target = &entry; + break; + } + } + + if (!target) + return CancellationCode::NotFound; + + if (target->descriptor.status != MergeTreePartitionExportTask::Status::PENDING) + return CancellationCode::CancelCannotBeSent; + + target->descriptor.status = MergeTreePartitionExportTask::Status::KILLED; + descriptor_json = target->descriptor.toJsonString(); + } + + persist(transaction_id, descriptor_json); + + /// Cancel any in-flight part exports for this transaction. Their completion callbacks will see + /// the KILLED status and skip further bookkeeping. + storage.killExportPart(transaction_id); + + return CancellationCode::CancelSent; +} + +std::vector MergeTreePartitionExportScheduler::getInfo() const +{ + std::vector result; + std::lock_guard lock(mutex); + result.reserve(tasks.size()); + for (const auto & [key, entry] : tasks) + { + const auto & descriptor = entry.descriptor; + PartitionExportInfo info; + info.source_database = descriptor.source_database; + info.source_table = descriptor.source_table; + info.destination_database = descriptor.destination_database; + info.destination_table = descriptor.destination_table; + info.create_time = descriptor.create_time; + info.partition_id = descriptor.partition_id; + info.transaction_id = descriptor.transaction_id; + info.query_id = descriptor.query_id; + info.parts = descriptor.partNames(); + info.parts_count = descriptor.partsCount(); + info.parts_to_do = descriptor.partsToDo(); + info.status = String(magic_enum::enum_name(descriptor.status)); + info.last_exception_message = descriptor.last_exception.message; + info.last_exception_part = descriptor.last_exception.part; + info.last_exception_time = descriptor.last_exception.time; + info.exception_count = descriptor.last_exception.count; + result.push_back(std::move(info)); + } + return result; +} + +void MergeTreePartitionExportScheduler::run() +{ + const auto available_move_executors = storage.background_moves_assignee.getAvailableMoveExecutors(); + if (available_move_executors == 0) + return; + + if (storage.parts_mover.moves_blocker.isCancelled()) + return; + + /// Respect the background memory soft-limit like the per-part export path does. + if (!canEnqueueBackgroundTask()) + return; + + std::vector> parts_to_schedule; /// (transaction_id, part_name) + std::vector tasks_to_commit; + + { + std::lock_guard lock(mutex); + size_t scheduled = 0; + for (auto & [key, entry] : tasks) + { + auto & descriptor = entry.descriptor; + if (descriptor.status != MergeTreePartitionExportTask::Status::PENDING) + continue; + + if (descriptor.allPartsDone()) + { + /// All parts exported: commit (or retry a previously-failed commit). Guard with the + /// committing flag so a concurrent completion callback does not commit in parallel. + if (!entry.committing) + { + entry.committing = true; + tasks_to_commit.push_back(descriptor.transaction_id); + } + continue; + } + + for (const auto & part : descriptor.parts) + { + if (scheduled >= available_move_executors) + break; + if (part.done || entry.in_flight_parts.contains(part.part_name)) + continue; + + entry.in_flight_parts.insert(part.part_name); + parts_to_schedule.emplace_back(descriptor.transaction_id, part.part_name); + ++scheduled; + } + + if (scheduled >= available_move_executors) + break; + } + } + + for (const auto & [transaction_id, part_name] : parts_to_schedule) + scheduleOnePart(transaction_id, part_name); + + for (const auto & transaction_id : tasks_to_commit) + tryCommit(transaction_id); +} + +void MergeTreePartitionExportScheduler::scheduleOnePart(const String & transaction_id, const String & part_name) +{ + MergeTreePartitionExportTask descriptor_copy; + { + std::lock_guard lock(mutex); + TaskEntry * entry = nullptr; + for (auto & [key, candidate] : tasks) + if (candidate.descriptor.transaction_id == transaction_id) + { + entry = &candidate; + break; + } + if (!entry) + return; + descriptor_copy = entry->descriptor; + } + + const StorageID destination_storage_id{descriptor_copy.destination_database, descriptor_copy.destination_table}; + + auto release_in_flight = [this, transaction_id, part_name]() + { + std::lock_guard lock(mutex); + for (auto & [key, entry] : tasks) + if (entry.descriptor.transaction_id == transaction_id) + { + entry.in_flight_parts.erase(part_name); + break; + } + }; + + try + { + auto context = ExportPartitionUtils::getContextCopyWithTaskSettings(storage.getContext(), descriptor_copy); + + LOG_INFO(storage.log, "ExportPartition: scheduling part export {} for task {}", part_name, transaction_id); + + storage.exportPartToTable( + part_name, + destination_storage_id, + transaction_id, + context, + descriptor_copy.iceberg_metadata_json, + /*allow_outdated_parts*/ true, + [this, transaction_id, part_name](MergeTreePartExportManifest::CompletionCallbackResult result) + { + handlePartCompletion(transaction_id, part_name, result); + }); + } + catch (...) + { + tryLogCurrentException(storage.log, __PRETTY_FUNCTION__); + /// Dispatch failed (e.g. destination missing, executor busy). Release the in-flight marker + /// so the part becomes eligible again on the next tick. + release_in_flight(); + } +} + +void MergeTreePartitionExportScheduler::handlePartCompletion( + const String & transaction_id, const String & part_name, const MergeTreePartExportManifest::CompletionCallbackResult & result) +{ + bool ready_to_commit = false; + String descriptor_json; + bool should_persist = false; + + { + std::lock_guard lock(mutex); + TaskEntry * entry = nullptr; + for (auto & [key, candidate] : tasks) + if (candidate.descriptor.transaction_id == transaction_id) + { + entry = &candidate; + break; + } + if (!entry) + return; + + entry->in_flight_parts.erase(part_name); + + auto & descriptor = entry->descriptor; + + /// Task already terminal (KILLED / FAILED / COMPLETED): ignore late completions. + if (descriptor.status != MergeTreePartitionExportTask::Status::PENDING) + return; + + if (result.success) + { + if (auto * part = descriptor.findPart(part_name)) + { + part->done = true; + part->paths_in_destination = result.relative_paths_in_destination_storage; + } + should_persist = true; + + if (descriptor.allPartsDone() && !entry->committing) + { + entry->committing = true; + ready_to_commit = true; + } + } + else + { + /// A cancelled export (KILL, or SYSTEM STOP MOVES) is not a real failure: leave the part + /// pending. If it was a KILL the status is already handled above; otherwise the next tick + /// retries it. + if (result.exception && result.exception->code() == ErrorCodes::QUERY_WAS_CANCELLED) + return; + + descriptor.last_exception.message = result.exception ? result.exception->message() : "Unknown export failure"; + descriptor.last_exception.part = part_name; + descriptor.last_exception.time = time(nullptr); + descriptor.last_exception.count += 1; + should_persist = true; + + if (result.exception && ExportPartitionUtils::isNonRetryableExportError(result.exception->code())) + { + descriptor.status = MergeTreePartitionExportTask::Status::FAILED; + LOG_WARNING(storage.log, "ExportPartition: task {} failed on part {} with non-retryable error (code {})", + transaction_id, part_name, result.exception->code()); + } + else + { + LOG_INFO(storage.log, "ExportPartition: task {} part {} failed with retryable error, will retry", + transaction_id, part_name); + } + } + + if (should_persist) + descriptor_json = descriptor.toJsonString(); + } + + if (should_persist) + persist(transaction_id, descriptor_json); + + if (ready_to_commit) + tryCommit(transaction_id); +} + +void MergeTreePartitionExportScheduler::tryCommit(const String & transaction_id) +{ + MergeTreePartitionExportTask descriptor_copy; + { + std::lock_guard lock(mutex); + TaskEntry * entry = nullptr; + for (auto & [key, candidate] : tasks) + if (candidate.descriptor.transaction_id == transaction_id) + { + entry = &candidate; + break; + } + if (!entry) + return; + + if (entry->descriptor.status != MergeTreePartitionExportTask::Status::PENDING || !entry->descriptor.allPartsDone()) + { + entry->committing = false; + return; + } + + entry->committing = true; + descriptor_copy = entry->descriptor; + } + + const StorageID destination_storage_id{descriptor_copy.destination_database, descriptor_copy.destination_table}; + + bool success = false; + std::optional failure; + try + { + auto destination_storage = DatabaseCatalog::instance().tryGetTable(destination_storage_id, storage.getContext()); + if (!destination_storage) + throw Exception(ErrorCodes::UNKNOWN_TABLE, "Destination table {} not found for export commit", + destination_storage_id.getNameForLogs()); + + const auto exported_paths = descriptor_copy.collectExportedPaths(); + if (exported_paths.empty()) + throw Exception(ErrorCodes::CORRUPTED_DATA, + "No exported paths found for export {}, will not commit. This might be a bug", transaction_id); + + auto context = ExportPartitionUtils::getContextCopyWithTaskSettings(storage.getContext(), descriptor_copy); + + LOG_INFO(storage.log, "ExportPartition: all parts exported for task {}, committing", transaction_id); + + ExportPartitionUtils::commitExportedPaths( + descriptor_copy.transaction_id, + descriptor_copy.partition_id, + descriptor_copy.iceberg_metadata_json, + descriptor_copy.write_full_path_in_iceberg_metadata, + exported_paths, + destination_storage, + storage, + context); + + success = true; + } + catch (const Exception & e) + { + failure = e; + LOG_WARNING(storage.log, "ExportPartition: commit for task {} failed: {}", transaction_id, e.message()); + } + + String descriptor_json; + { + std::lock_guard lock(mutex); + TaskEntry * entry = nullptr; + for (auto & [key, candidate] : tasks) + if (candidate.descriptor.transaction_id == transaction_id) + { + entry = &candidate; + break; + } + if (!entry) + return; + + entry->committing = false; + auto & descriptor = entry->descriptor; + + /// A concurrent KILL may have won the race while we were committing. + if (descriptor.status != MergeTreePartitionExportTask::Status::PENDING) + { + descriptor_json = descriptor.toJsonString(); + } + else if (success) + { + descriptor.status = MergeTreePartitionExportTask::Status::COMPLETED; + descriptor_json = descriptor.toJsonString(); + } + else + { + descriptor.last_exception.message = failure ? failure->message() : "Unknown commit failure"; + descriptor.last_exception.part = ""; + descriptor.last_exception.time = time(nullptr); + descriptor.last_exception.count += 1; + + if (failure && ExportPartitionUtils::isNonRetryableExportError(failure->code())) + descriptor.status = MergeTreePartitionExportTask::Status::FAILED; + /// Otherwise leave PENDING: run() will retry the commit on the next tick. + + descriptor_json = descriptor.toJsonString(); + } + } + + persist(transaction_id, descriptor_json); +} + +void MergeTreePartitionExportScheduler::persist(const String & transaction_id, const String & descriptor_json) +{ + std::lock_guard file_lock(persist_mutex); + + auto disk = storage.getDisks().front(); + const auto directory = getExportsRelativePath(); + disk->createDirectories(directory); + + const auto final_path = fs::path(directory) / (transaction_id + ".json"); + const auto tmp_path = fs::path(directory) / (transaction_id + ".json.tmp"); + + { + auto out = disk->writeFile(tmp_path, DBMS_DEFAULT_BUFFER_SIZE, WriteMode::Rewrite, storage.getContext()->getWriteSettings()); + writeString(descriptor_json, *out); + out->finalize(); + } + + disk->replaceFile(tmp_path, final_path); +} + +void MergeTreePartitionExportScheduler::removeTaskFile(const String & transaction_id) +{ + std::lock_guard file_lock(persist_mutex); + auto disk = storage.getDisks().front(); + const auto final_path = fs::path(getExportsRelativePath()) / (transaction_id + ".json"); + disk->removeFileIfExists(final_path); +} + +void MergeTreePartitionExportScheduler::loadFromDisk() +{ + auto disk = storage.getDisks().front(); + const auto directory = getExportsRelativePath(); + + if (!disk->existsDirectory(directory)) + return; + + std::lock_guard lock(mutex); + + for (auto it = disk->iterateDirectory(directory); it->isValid(); it->next()) + { + const auto file_name = it->name(); + /// Skip stale temporary files left by an interrupted write. + if (!endsWith(file_name, ".json")) + continue; + + try + { + auto buf = disk->readFile(fs::path(directory) / file_name, getReadSettings()); + String content; + readStringUntilEOF(content, *buf); + + auto descriptor = MergeTreePartitionExportTask::fromJsonString(content); + const auto composite_key = compositeKey(descriptor.partition_id, descriptor.destination_database, descriptor.destination_table); + + TaskEntry entry; + + /// Re-pin the not-yet-exported parts of a resumable task so background merges do not + /// remove them before the export finishes. + if (descriptor.status == MergeTreePartitionExportTask::Status::PENDING) + { + for (const auto & part : descriptor.parts) + { + if (part.done) + continue; + if (auto data_part = storage.getPartIfExists( + part.part_name, {MergeTreeDataPartState::Active, MergeTreeDataPartState::Outdated})) + entry.part_references.push_back(data_part); + } + } + + entry.descriptor = std::move(descriptor); + tasks.emplace(composite_key, std::move(entry)); + + LOG_INFO(storage.log, "ExportPartition: loaded export task from disk (key {})", composite_key); + } + catch (...) + { + tryLogCurrentException(storage.log, "Failed to load partition export descriptor " + file_name); + } + } +} + +} diff --git a/src/Storages/MergeTree/MergeTreePartitionExportScheduler.h b/src/Storages/MergeTree/MergeTreePartitionExportScheduler.h new file mode 100644 index 000000000000..15a6adbcad07 --- /dev/null +++ b/src/Storages/MergeTree/MergeTreePartitionExportScheduler.h @@ -0,0 +1,111 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace DB +{ + +class StorageMergeTree; + +/// Read-model row for `system.partition_exports`. Populated from the scheduler's in-memory registry +/// (no disk I/O). Mirrors the useful subset of the replicated info, without the replica-specific +/// columns that make no sense for a single-node table. +struct PartitionExportInfo +{ + String source_database; + String source_table; + String destination_database; + String destination_table; + time_t create_time = 0; + String partition_id; + String transaction_id; + String query_id; + std::vector parts; + size_t parts_count = 0; + size_t parts_to_do = 0; + String status; + String last_exception_message; + String last_exception_part; + time_t last_exception_time = 0; + size_t exception_count = 0; +}; + +/// Local (ZooKeeper-free) coordinator for `EXPORT PARTITION` on a plain `MergeTree` table. +/// +/// A single node owns the whole task, so there are no locks, no cross-replica scheduling, and no +/// races: the registry below is the live source of truth, backed by one JSON descriptor file per +/// task on disk (so tasks resume after a restart). The registry is guarded by a plain mutex; all +/// slow I/O (disk persistence, the object-storage / Iceberg commit, and scheduling part exports) +/// is performed outside the lock. +class MergeTreePartitionExportScheduler +{ +public: + explicit MergeTreePartitionExportScheduler(StorageMergeTree & storage_); + + using DataPartPtr = MergeTreePartExportManifest::DataPartPtr; + + /// Registers and persists a new task, then triggers the scheduler. Throws + /// EXPORT_PARTITION_ALREADY_EXPORTED if a task with the same (partition, destination) key + /// already exists and `force` is false; otherwise the previous task (and its file) is replaced. + void addTask(MergeTreePartitionExportTask descriptor, std::vector part_references, bool force); + + /// Best-effort early duplicate check used before doing the heavy request-time validation. + /// The authoritative atomic check happens inside addTask. + void throwIfAlreadyExported(const String & composite_key, bool force) const; + + /// Cancels a PENDING task: flips its status to KILLED, persists, and cancels any in-flight + /// part-export tasks for the transaction. Returns a CancellationCode for the KILL query. + CancellationCode kill(const String & transaction_id); + + /// Snapshot of every tracked task for system.partition_exports. Briefly locks the registry. + std::vector getInfo() const; + + /// Scheduler tick: schedule pending parts of PENDING tasks and commit tasks whose parts are all + /// exported. Invoked periodically from the storage's schedule-pool task. + void run(); + + /// Reloads persisted descriptors from disk and re-pins the parts of PENDING tasks. Called once + /// during table startup, before background merges can remove parts. + void loadFromDisk(); + + static String compositeKey(const String & partition_id, const String & destination_database, const String & destination_table); + +private: + StorageMergeTree & storage; + + struct TaskEntry + { + MergeTreePartitionExportTask descriptor; + /// Pins the source parts so they are not physically removed before the export finishes. + std::vector part_references; + /// Parts currently scheduled on the background move executor (avoids double scheduling). + std::unordered_set in_flight_parts; + /// True while a commit attempt is in progress, so run() and completion callbacks do not + /// drive commitExportPartitionTransaction concurrently for the same task. + bool committing = false; + }; + + mutable std::mutex mutex; + std::map tasks; + + /// Serializes on-disk descriptor writes so two concurrent completion callbacks cannot interleave. + std::mutex persist_mutex; + + void scheduleOnePart(const String & transaction_id, const String & part_name); + void handlePartCompletion(const String & transaction_id, const String & part_name, const MergeTreePartExportManifest::CompletionCallbackResult & result); + void tryCommit(const String & transaction_id); + + /// Serialize `descriptor` to its on-disk file (atomic tmp + replace). No registry lock is held. + void persist(const String & transaction_id, const String & descriptor_json); + void removeTaskFile(const String & transaction_id); + + String getExportsRelativePath() const; +}; + +} diff --git a/src/Storages/MergeTree/MergeTreePartitionExportTask.h b/src/Storages/MergeTree/MergeTreePartitionExportTask.h new file mode 100644 index 000000000000..ea0f18731c82 --- /dev/null +++ b/src/Storages/MergeTree/MergeTreePartitionExportTask.h @@ -0,0 +1,250 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace DB +{ + +/// On-disk descriptor for a plain (non-replicated) `MergeTree` partition export task. +/// +/// Unlike the replicated variant, there is a single node and no cross-replica coordination, +/// so this descriptor is the whole source of truth. It is persisted as one JSON file per task +/// under the table data directory and rewritten on every state transition (part done, status +/// change). The scheduler keeps an in-memory copy guarded by a mutex; the descriptor is small +/// and self-contained so it can be reloaded verbatim after a server restart. +struct MergeTreePartitionExportTask +{ + using FileAlreadyExistsPolicy = MergeTreePartExportManifest::FileAlreadyExistsPolicy; + + enum class Status + { + PENDING, + COMPLETED, + FAILED, + KILLED, + }; + + struct PartProgress + { + String part_name; + bool done = false; + std::vector paths_in_destination; + }; + + /// Best-effort record of the most recent failure observed for this task. `count` is a running + /// total of failures and matches the semantics documented for `system.partition_exports`. + struct LastException + { + String message; + String part; /// empty for task-level exceptions (commit failure) + time_t time = 0; + size_t count = 0; + }; + + /// Identity + String transaction_id; + String query_id; + String partition_id; + String source_database; + String source_table; + String destination_database; + String destination_table; + time_t create_time = 0; + + /// Work + progress + std::vector parts; + Status status = Status::PENDING; + LastException last_exception; + + /// Export settings. Field names are intentionally identical to + /// `ExportReplicatedMergeTreePartitionManifest` so that + /// `ExportPartitionUtils::getContextCopyWithTaskSettings` works for both descriptors. + size_t max_threads = 0; + bool parallel_formatting = false; + bool parquet_parallel_encoding = false; + size_t max_bytes_per_file = 0; + size_t max_rows_per_file = 0; + FileAlreadyExistsPolicy file_already_exists_policy = FileAlreadyExistsPolicy::skip; + String filename_pattern; + bool write_full_path_in_iceberg_metadata = false; + bool allow_lossy_cast = false; + String iceberg_metadata_json; + String parquet_compression_method; + UInt64 output_format_compression_level = 0; + UInt64 parquet_row_group_size = 0; + UInt64 parquet_row_group_size_bytes = 0; + + size_t partsCount() const { return parts.size(); } + + size_t partsToDo() const + { + size_t to_do = 0; + for (const auto & part : parts) + if (!part.done) + ++to_do; + return to_do; + } + + bool allPartsDone() const + { + for (const auto & part : parts) + if (!part.done) + return false; + return true; + } + + std::vector partNames() const + { + std::vector names; + names.reserve(parts.size()); + for (const auto & part : parts) + names.push_back(part.part_name); + return names; + } + + /// Flattened list of every destination path produced by the already-exported parts. + std::vector collectExportedPaths() const + { + std::vector paths; + for (const auto & part : parts) + for (const auto & path : part.paths_in_destination) + paths.push_back(path); + return paths; + } + + PartProgress * findPart(const String & part_name) + { + for (auto & part : parts) + if (part.part_name == part_name) + return ∂ + return nullptr; + } + + std::string toJsonString() const + { + Poco::JSON::Object json; + json.set("transaction_id", transaction_id); + json.set("query_id", query_id); + json.set("partition_id", partition_id); + json.set("source_database", source_database); + json.set("source_table", source_table); + json.set("destination_database", destination_database); + json.set("destination_table", destination_table); + json.set("create_time", create_time); + json.set("status", String(magic_enum::enum_name(status))); + + Poco::JSON::Array::Ptr parts_array = new Poco::JSON::Array(); + for (const auto & part : parts) + { + Poco::JSON::Object::Ptr part_object = new Poco::JSON::Object(); + part_object->set("part_name", part.part_name); + part_object->set("done", part.done); + Poco::JSON::Array::Ptr paths_array = new Poco::JSON::Array(); + for (const auto & path : part.paths_in_destination) + paths_array->add(path); + part_object->set("paths_in_destination", paths_array); + parts_array->add(part_object); + } + json.set("parts", parts_array); + + Poco::JSON::Object::Ptr exception_object = new Poco::JSON::Object(); + exception_object->set("message", last_exception.message); + exception_object->set("part", last_exception.part); + exception_object->set("time", last_exception.time); + exception_object->set("count", last_exception.count); + json.set("last_exception", exception_object); + + json.set("max_threads", max_threads); + json.set("parallel_formatting", parallel_formatting); + json.set("parquet_parallel_encoding", parquet_parallel_encoding); + json.set("max_bytes_per_file", max_bytes_per_file); + json.set("max_rows_per_file", max_rows_per_file); + json.set("file_already_exists_policy", String(magic_enum::enum_name(file_already_exists_policy))); + json.set("filename_pattern", filename_pattern); + json.set("write_full_path_in_iceberg_metadata", write_full_path_in_iceberg_metadata); + json.set("allow_lossy_cast", allow_lossy_cast); + if (!iceberg_metadata_json.empty()) + json.set("iceberg_metadata_json", iceberg_metadata_json); + json.set("parquet_compression_method", parquet_compression_method); + json.set("output_format_compression_level", output_format_compression_level); + json.set("parquet_row_group_size", parquet_row_group_size); + json.set("parquet_row_group_size_bytes", parquet_row_group_size_bytes); + + std::ostringstream oss; // STYLE_CHECK_ALLOW_STD_STRING_STREAM + oss.exceptions(std::ios::failbit); + Poco::JSON::Stringifier::stringify(json, oss); + return oss.str(); + } + + static MergeTreePartitionExportTask fromJsonString(const std::string & json_string) + { + Poco::JSON::Parser parser; + auto json = parser.parse(json_string).extract(); + chassert(json); + + MergeTreePartitionExportTask task; + task.transaction_id = json->getValue("transaction_id"); + task.query_id = json->getValue("query_id"); + task.partition_id = json->getValue("partition_id"); + task.source_database = json->getValue("source_database"); + task.source_table = json->getValue("source_table"); + task.destination_database = json->getValue("destination_database"); + task.destination_table = json->getValue("destination_table"); + task.create_time = json->getValue("create_time"); + + if (const auto status = magic_enum::enum_cast(json->getValue("status"))) + task.status = status.value(); + + const auto parts_array = json->getArray("parts"); + for (size_t i = 0; i < parts_array->size(); ++i) + { + const auto part_object = parts_array->getObject(static_cast(i)); + PartProgress part; + part.part_name = part_object->getValue("part_name"); + part.done = part_object->getValue("done"); + const auto paths_array = part_object->getArray("paths_in_destination"); + for (size_t j = 0; j < paths_array->size(); ++j) + part.paths_in_destination.push_back(paths_array->getElement(static_cast(j))); + task.parts.push_back(std::move(part)); + } + + if (json->has("last_exception")) + { + const auto exception_object = json->getObject("last_exception"); + task.last_exception.message = exception_object->getValue("message"); + task.last_exception.part = exception_object->getValue("part"); + task.last_exception.time = exception_object->getValue("time"); + task.last_exception.count = exception_object->getValue("count"); + } + + task.max_threads = json->getValue("max_threads"); + task.parallel_formatting = json->getValue("parallel_formatting"); + task.parquet_parallel_encoding = json->getValue("parquet_parallel_encoding"); + task.max_bytes_per_file = json->getValue("max_bytes_per_file"); + task.max_rows_per_file = json->getValue("max_rows_per_file"); + + if (const auto policy = magic_enum::enum_cast(json->getValue("file_already_exists_policy"))) + task.file_already_exists_policy = policy.value(); + + task.filename_pattern = json->getValue("filename_pattern"); + task.write_full_path_in_iceberg_metadata = json->getValue("write_full_path_in_iceberg_metadata"); + task.allow_lossy_cast = json->getValue("allow_lossy_cast"); + if (json->has("iceberg_metadata_json")) + task.iceberg_metadata_json = json->getValue("iceberg_metadata_json"); + task.parquet_compression_method = json->getValue("parquet_compression_method"); + task.output_format_compression_level = json->getValue("output_format_compression_level"); + task.parquet_row_group_size = json->getValue("parquet_row_group_size"); + task.parquet_row_group_size_bytes = json->getValue("parquet_row_group_size_bytes"); + + return task; + } +}; + +} diff --git a/src/Storages/StorageMergeTree.cpp b/src/Storages/StorageMergeTree.cpp index fb65722ae05e..e5294f7a074f 100644 --- a/src/Storages/StorageMergeTree.cpp +++ b/src/Storages/StorageMergeTree.cpp @@ -7,7 +7,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -22,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -43,6 +46,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -93,6 +99,30 @@ namespace Setting extern const SettingsBool throw_on_unsupported_query_inside_transaction; extern const SettingsUInt64 max_parts_to_move; extern const SettingsUpdateParallelMode update_parallel_mode; + extern const SettingsBool export_merge_tree_partition_force_export; + extern const SettingsBool output_format_parallel_formatting; + extern const SettingsBool output_format_parquet_parallel_encoding; + extern const SettingsParquetCompression output_format_parquet_compression_method; + extern const SettingsUInt64 output_format_compression_level; + extern const SettingsUInt64 output_format_parquet_row_group_size; + extern const SettingsUInt64 output_format_parquet_row_group_size_bytes; + extern const SettingsMaxThreads max_threads; + extern const SettingsMergeTreePartExportFileAlreadyExistsPolicy export_merge_tree_part_file_already_exists_policy; + extern const SettingsUInt64 export_merge_tree_part_max_bytes_per_file; + extern const SettingsUInt64 export_merge_tree_part_max_rows_per_file; + extern const SettingsBool export_merge_tree_part_throw_on_pending_mutations; + extern const SettingsBool export_merge_tree_part_throw_on_pending_patch_parts; + extern const SettingsBool export_merge_tree_part_allow_lossy_cast; + extern const SettingsExportPartitionAllOnError export_merge_tree_partition_all_on_error; + extern const SettingsString export_merge_tree_part_filename_pattern; + extern const SettingsBool write_full_path_in_iceberg_metadata; + extern const SettingsUInt64 iceberg_insert_max_bytes_in_data_file; + extern const SettingsUInt64 iceberg_insert_max_rows_in_data_file; +} + +namespace ServerSetting +{ + extern const ServerSettingsBool allow_experimental_export_merge_tree_partition; } namespace MergeTreeSetting @@ -134,6 +164,9 @@ namespace ErrorCodes extern const int PART_IS_LOCKED; extern const int PART_IS_TEMPORARILY_LOCKED; extern const int INCOMPATIBLE_COLUMNS; + extern const int EXPORT_PARTITION_ALREADY_EXPORTED; + extern const int PARTITION_EXPORT_FAILED; + extern const int PENDING_MUTATIONS_NOT_ALLOWED; } namespace ActionLocks @@ -212,6 +245,19 @@ StorageMergeTree::StorageMergeTree( loadMutations(); loadDeduplicationLog(); prewarmCaches(getActivePartsLoadingThreadPool().get(), getCachesToPrewarm(0)); + + if (getContext()->getServerSettings()[ServerSetting::allow_experimental_export_merge_tree_partition]) + { + partition_export_scheduler = std::make_shared(*this); + + partition_export_task = getContext()->getSchedulePool().createTask( + getStorageID(), + getStorageID().getFullTableName() + " (StorageMergeTree::partition_export_task)", + [this] { partitionExportTask(); }); + + /// Activated in startup(); deactivated during shutdown. + partition_export_task->deactivate(); + } } @@ -231,6 +277,14 @@ void StorageMergeTree::startup() try { + /// Reload persisted partition-export tasks (and re-pin their parts) before background merges + /// can start removing parts, then activate the scheduler task so PENDING tasks resume. + if (partition_export_scheduler) + { + partition_export_scheduler->loadFromDisk(); + partition_export_task->activateAndSchedule(); + } + cleanup_thread.start(); background_operations_assignee.start(); startBackgroundMoves(); @@ -266,6 +320,9 @@ void StorageMergeTree::flushAndPrepareForShutdown() merger_mutator.merges_blocker.cancelForever(); parts_mover.moves_blocker.cancelForever(); + if (partition_export_task) + partition_export_task->deactivate(); + background_operations_assignee.finish(); background_moves_assignee.finish(); @@ -3185,4 +3242,222 @@ CommittingBlocksSet StorageMergeTree::getCommittingBlocks() const std::lock_guard lock(committing_blocks_mutex); return committing_blocks; } + +void StorageMergeTree::exportPartitionToTable(const PartitionCommand & command, ContextPtr query_context) +{ + if (!query_context->getServerSettings()[ServerSetting::allow_experimental_export_merge_tree_partition]) + throw Exception(ErrorCodes::SUPPORT_IS_DISABLED, + "Exporting merge tree partition is experimental. Set the server setting `allow_experimental_export_merge_tree_partition` to enable it.\n" + "If you are exporting to an Apache Iceberg table, you also need to enable the setting `allow_insert_into_iceberg`."); + + /// The scheduler is created in the constructor whenever the server setting above is enabled, so + /// this should always hold here. + if (!partition_export_scheduler) + throw Exception(ErrorCodes::SUPPORT_IS_DISABLED, + "Partition export is not initialized for table {}", getStorageID().getNameForLogs()); + + /// EXPORT PARTITION ALL: expand into one sub-call per active partition id. + /// Failure handling is controlled by `export_merge_tree_partition_all_on_error`. + if (const auto * partition_ast = command.partition->as(); partition_ast && partition_ast->all) + { + auto partition_id_set = getAllPartitionIds(); + if (partition_id_set.empty()) + throw Exception(ErrorCodes::BAD_ARGUMENTS, + "Table {} has no active partitions to export", getStorageID().getNameForLogs()); + + std::vector partition_ids(partition_id_set.begin(), partition_id_set.end()); + std::sort(partition_ids.begin(), partition_ids.end()); + + const auto & on_error_setting = query_context->getSettingsRef()[Setting::export_merge_tree_partition_all_on_error]; + const ExportPartitionAllOnError on_error = on_error_setting.value; + + LOG_INFO(log, "EXPORT PARTITION ALL: scheduling export for {} partitions, on_error={}", + partition_ids.size(), on_error_setting.toString()); + + std::vector> failures; /// (partition_id, message) + size_t skipped_conflicts = 0; + + for (const auto & partition_id : partition_ids) + { + PartitionCommand sub = command; + auto synthetic = make_intrusive(); + synthetic->setPartitionID(make_intrusive(partition_id)); + sub.partition = synthetic; + + try + { + exportPartitionToTable(sub, query_context); + } + catch (const Exception & e) + { + switch (on_error) + { + case ExportPartitionAllOnError::throw_first: + throw; + case ExportPartitionAllOnError::skip_conflicts: + if (e.code() == ErrorCodes::EXPORT_PARTITION_ALREADY_EXPORTED) + { + ++skipped_conflicts; + LOG_INFO(log, "EXPORT PARTITION ALL: skipping partition {} (already exported): {}", + partition_id, e.message()); + break; + } + throw; + case ExportPartitionAllOnError::collect: + LOG_WARNING(log, "EXPORT PARTITION ALL: partition {} failed: {}", partition_id, e.message()); + failures.emplace_back(partition_id, e.message()); + break; + } + } + } + + if (!failures.empty()) + { + String aggregated = fmt::format( + "EXPORT PARTITION ALL: {}/{} partitions failed to schedule. Per-partition errors:", + failures.size(), partition_ids.size()); + for (const auto & [pid, msg] : failures) + aggregated += fmt::format("\n {}: {}", pid, msg); + throw Exception(ErrorCodes::PARTITION_EXPORT_FAILED, "{}", aggregated); + } + + if (skipped_conflicts > 0) + LOG_INFO(log, "EXPORT PARTITION ALL: skipped {} partitions due to existing exports", skipped_conflicts); + + return; + } + + const auto dest_database = query_context->resolveDatabase(command.to_database); + const auto dest_table = command.to_table; + const auto dest_storage_id = StorageID(dest_database, dest_table); + auto dest_storage = DatabaseCatalog::instance().getTable({dest_database, dest_table}, query_context); + + auto src_snapshot = getInMemoryMetadataPtr(); + + /// Validate the destination (not self, import support, positionally castable schema, matching + /// partition key / Iceberg partition spec) and capture the destination metadata.json (empty for + /// non-data-lake destinations). + const auto destination_iceberg_metadata_json = ExportPartitionUtils::extractDestinationIcebergMetadataJson( + src_snapshot, getStorageID(), dest_storage, query_context); + + const String partition_id = getPartitionIDFromQuery(command.partition, query_context); + const auto composite_key = MergeTreePartitionExportScheduler::compositeKey(partition_id, dest_database, dest_table); + + const bool force = query_context->getSettingsRef()[Setting::export_merge_tree_partition_force_export]; + + /// Fail fast on a duplicate before doing part collection and mutation checks. The authoritative + /// atomic check happens again inside addTask. + partition_export_scheduler->throwIfAlreadyExported(composite_key, force); + + DataPartsVector parts; + { + auto data_parts_lock = lockParts(); + parts = getDataPartsVectorInPartitionForInternalUsage(MergeTreeDataPartState::Active, partition_id, data_parts_lock); + } + + if (parts.empty()) + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Partition {} doesn't exist", partition_id); + + const bool throw_on_pending_mutations = query_context->getSettingsRef()[Setting::export_merge_tree_part_throw_on_pending_mutations]; + const bool throw_on_pending_patch_parts = query_context->getSettingsRef()[Setting::export_merge_tree_part_throw_on_pending_patch_parts]; + + MergeTreeData::IMutationsSnapshot::Params mutations_snapshot_params + { + .metadata_version = getInMemoryMetadataPtr()->getMetadataVersion(), + .min_part_metadata_version = MergeTreeData::getPartsSnapshotInfo(parts).min_metadata_version, + .need_data_mutations = throw_on_pending_mutations, + .need_alter_mutations = throw_on_pending_mutations || throw_on_pending_patch_parts, + .need_patch_parts = throw_on_pending_patch_parts, + }; + + const auto mutations_snapshot = getMutationsSnapshot(mutations_snapshot_params); + + for (const auto & part : parts) + { + const auto alter_conversions = getAlterConversionsForPart(part, mutations_snapshot, query_context); + + /// re-check `throw_on_pending_mutations` because `pending_mutations` might have been filled due to `throw_on_pending_patch_parts` + if (alter_conversions->hasMutations() && throw_on_pending_mutations) + throw Exception(ErrorCodes::PENDING_MUTATIONS_NOT_ALLOWED, + "Partition {} can not be exported because the part {} has pending mutations. Either wait for the mutations to be applied or set `export_merge_tree_part_throw_on_pending_mutations` to false", + partition_id, part->name); + + if (alter_conversions->hasPatches()) + throw Exception(ErrorCodes::PENDING_MUTATIONS_NOT_ALLOWED, + "Partition {} can not be exported because the part {} has pending patch parts. Either wait for the patch parts to be applied or set `export_merge_tree_part_throw_on_pending_patch_parts` to false", + partition_id, part->name); + } + + MergeTreePartitionExportTask descriptor; + descriptor.transaction_id = generateSnowflakeIDString(); + descriptor.query_id = query_context->getCurrentQueryId(); + descriptor.partition_id = partition_id; + descriptor.source_database = getStorageID().database_name; + descriptor.source_table = getStorageID().table_name; + descriptor.destination_database = dest_database; + descriptor.destination_table = dest_table; + descriptor.create_time = time(nullptr); + descriptor.status = MergeTreePartitionExportTask::Status::PENDING; + + for (const auto & part : parts) + descriptor.parts.push_back({part->name, /*done*/ false, /*paths*/ {}}); + + descriptor.max_threads = query_context->getSettingsRef()[Setting::max_threads]; + descriptor.parallel_formatting = query_context->getSettingsRef()[Setting::output_format_parallel_formatting]; + descriptor.parquet_parallel_encoding = query_context->getSettingsRef()[Setting::output_format_parquet_parallel_encoding]; + descriptor.parquet_compression_method = query_context->getSettingsRef()[Setting::output_format_parquet_compression_method].toString(); + descriptor.output_format_compression_level = query_context->getSettingsRef()[Setting::output_format_compression_level]; + descriptor.parquet_row_group_size = query_context->getSettingsRef()[Setting::output_format_parquet_row_group_size]; + descriptor.parquet_row_group_size_bytes = query_context->getSettingsRef()[Setting::output_format_parquet_row_group_size_bytes]; + descriptor.max_bytes_per_file = query_context->getSettingsRef()[Setting::export_merge_tree_part_max_bytes_per_file]; + descriptor.max_rows_per_file = query_context->getSettingsRef()[Setting::export_merge_tree_part_max_rows_per_file]; + descriptor.file_already_exists_policy = query_context->getSettingsRef()[Setting::export_merge_tree_part_file_already_exists_policy].value; + descriptor.filename_pattern = query_context->getSettingsRef()[Setting::export_merge_tree_part_filename_pattern].value; + descriptor.write_full_path_in_iceberg_metadata = query_context->getSettingsRef()[Setting::write_full_path_in_iceberg_metadata]; + descriptor.allow_lossy_cast = query_context->getSettingsRef()[Setting::export_merge_tree_part_allow_lossy_cast]; + + if (dest_storage->isDataLake()) + { + descriptor.iceberg_metadata_json = destination_iceberg_metadata_json; + descriptor.max_bytes_per_file = query_context->getSettingsRef()[Setting::iceberg_insert_max_bytes_in_data_file]; + descriptor.max_rows_per_file = query_context->getSettingsRef()[Setting::iceberg_insert_max_rows_in_data_file]; + } + + std::vector part_references(parts.begin(), parts.end()); + partition_export_scheduler->addTask(std::move(descriptor), std::move(part_references), force); +} + +CancellationCode StorageMergeTree::killExportPartition(const String & transaction_id) +{ + if (!partition_export_scheduler) + return CancellationCode::NotFound; + return partition_export_scheduler->kill(transaction_id); +} + +std::vector StorageMergeTree::getPartitionExportsInfo() const +{ + if (!partition_export_scheduler) + return {}; + return partition_export_scheduler->getInfo(); +} + +void StorageMergeTree::partitionExportTask() +{ + try + { + partition_export_scheduler->run(); + } + catch (...) + { + tryLogCurrentException(log, __PRETTY_FUNCTION__); + } + + partition_export_task->scheduleAfter(5000); +} + +void StorageMergeTree::triggerPartitionExportTask() +{ + if (partition_export_task) + partition_export_task->schedule(); +} } diff --git a/src/Storages/StorageMergeTree.h b/src/Storages/StorageMergeTree.h index 87ec88e418f5..a50f8a499834 100644 --- a/src/Storages/StorageMergeTree.h +++ b/src/Storages/StorageMergeTree.h @@ -1,10 +1,12 @@ #pragma once #include +#include #include #include #include #include +#include #include #include #include @@ -122,7 +124,28 @@ class StorageMergeTree final : public MergeTreeData MergeTreeDeduplicationLog * getDeduplicationLog() { return deduplication_log.get(); } + /// EXPORT PARTITION for a plain (non-replicated) MergeTree table. Coordinated locally by + /// `partition_export_scheduler`; the task descriptor is persisted on disk (no ZooKeeper). + void exportPartitionToTable(const PartitionCommand & command, ContextPtr query_context) override; + + CancellationCode killExportPartition(const String & transaction_id) override; + + /// Snapshot of local partition-export tasks for `system.partition_exports`. No disk I/O. + std::vector getPartitionExportsInfo() const; + private: + friend class MergeTreePartitionExportScheduler; + + /// Local coordinator + on-disk state for EXPORT PARTITION. Only created when the server setting + /// `allow_experimental_export_merge_tree_partition` is enabled. + std::shared_ptr partition_export_scheduler; + BackgroundSchedulePoolTaskHolder partition_export_task; + + /// Schedule-pool task body: drives partition_export_scheduler->run() periodically. + void partitionExportTask(); + /// Wakes up the partition-export schedule-pool task (no-op when the feature is disabled). + void triggerPartitionExportTask(); + /// Mutex and condvar for synchronous mutations wait std::mutex mutation_wait_mutex; diff --git a/src/Storages/StorageReplicatedMergeTree.cpp b/src/Storages/StorageReplicatedMergeTree.cpp index fe72425a36d5..418df7d0aa39 100644 --- a/src/Storages/StorageReplicatedMergeTree.cpp +++ b/src/Storages/StorageReplicatedMergeTree.cpp @@ -8386,33 +8386,13 @@ void StorageReplicatedMergeTree::exportPartitionToTable(const PartitionCommand & const auto dest_storage_id = StorageID(dest_database, dest_table); auto dest_storage = DatabaseCatalog::instance().getTable({dest_database, dest_table}, query_context); - if (dest_storage->getStorageID() == this->getStorageID()) - { - throw Exception(ErrorCodes::BAD_ARGUMENTS, "Exporting to the same table is not allowed"); - } - - if (!dest_storage->supportsImport(query_context)) - throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Destination storage {} does not support MergeTree parts or uses unsupported partitioning", dest_storage->getName()); - - auto query_to_string = [] (const ASTPtr & ast) - { - return ast ? ast->formatWithSecretsOneLine() : ""; - }; - auto src_snapshot = getInMemoryMetadataPtr(); - auto destination_snapshot = dest_storage->getInMemoryMetadataPtr(); - /// Positional CAST matching, like `INSERT INTO dest SELECT * FROM src`. - ExportPartitionUtils::verifyExportSchemaCastable( - src_snapshot, destination_snapshot, dest_storage->getStorageID(), query_context); - - /// Iceberg partition compatibility is checked below; here we only need the - /// partition-key ASTs to match (partition-column types follow the lossy-cast gate). - if (!dest_storage->isDataLake()) - { - if (query_to_string(src_snapshot->getPartitionKeyAST()) != query_to_string(destination_snapshot->getPartitionKeyAST())) - throw Exception(ErrorCodes::BAD_ARGUMENTS, "Tables have different partition key"); - } + /// Validate the destination (not self, import support, positionally castable schema, matching + /// partition key / Iceberg partition spec) and, for Iceberg destinations, capture the + /// destination metadata.json to persist in the manifest. + const auto destination_iceberg_metadata_json = ExportPartitionUtils::extractDestinationIcebergMetadataJson( + src_snapshot, getStorageID(), dest_storage, query_context); zkutil::ZooKeeperPtr zookeeper = getZooKeeperAndAssertNotReadonly(); @@ -8534,50 +8514,11 @@ void StorageReplicatedMergeTree::exportPartitionToTable(const PartitionCommand & if (dest_storage->isDataLake()) { -#if USE_AVRO - auto * object_storage = dynamic_cast(dest_storage.get()); - auto * object_storage_cluster = dynamic_cast(dest_storage.get()); - - /// in theory this should never happen, but just in case - if (!object_storage && !object_storage_cluster) - { - throw Exception(ErrorCodes::BAD_ARGUMENTS, "Destination storage {} is not a StorageObjectStorage", dest_storage->getName()); - } - - IcebergMetadata * iceberg_metadata = nullptr; - if (object_storage) - iceberg_metadata = dynamic_cast(object_storage->getExternalMetadata(query_context)); - else if (object_storage_cluster) - iceberg_metadata = dynamic_cast(object_storage_cluster->getExternalMetadata(query_context)); - if (!iceberg_metadata) - { - throw Exception(ErrorCodes::BAD_ARGUMENTS, "Destination storage {} is a data lake but not an iceberg table", dest_storage->getName()); - } - - if (!query_context->getSettingsRef()[Setting::allow_insert_into_iceberg]) - { - throw Exception(ErrorCodes::SUPPORT_IS_DISABLED, - "Iceberg writes are experimental. " - "To allow its usage, enable the setting `allow_insert_into_iceberg` on the initiator (query, session or profile) - replicas inherit it from the scheduled task."); - } - - const auto metadata_object = iceberg_metadata->getMetadataJSON(query_context); - - ExportPartitionUtils::verifyIcebergPartitionCompatibility( - metadata_object, - src_snapshot->getPartitionKeyAST()); - - std::ostringstream oss; // STYLE_CHECK_ALLOW_STD_STRING_STREAM - oss.exceptions(std::ios::failbit); - metadata_object->stringify(oss); - manifest.iceberg_metadata_json = oss.str(); - + /// Destination validation and metadata.json capture already happened up-front in + /// extractDestinationIcebergMetadataJson; here we only persist the results in the manifest. + manifest.iceberg_metadata_json = destination_iceberg_metadata_json; manifest.max_bytes_per_file = query_context->getSettingsRef()[Setting::iceberg_insert_max_bytes_in_data_file]; manifest.max_rows_per_file = query_context->getSettingsRef()[Setting::iceberg_insert_max_rows_in_data_file]; - -#else - throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Data lake export requires Avro support"); -#endif } ops.emplace_back(zkutil::makeCreateRequest( diff --git a/src/Storages/System/StorageSystemPartitionExports.cpp b/src/Storages/System/StorageSystemPartitionExports.cpp new file mode 100644 index 000000000000..caffe683a95a --- /dev/null +++ b/src/Storages/System/StorageSystemPartitionExports.cpp @@ -0,0 +1,144 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace DB +{ + +ColumnsDescription StorageSystemPartitionExports::getColumnsDescription() +{ + return ColumnsDescription + { + {"source_database", std::make_shared(), "Name of the source database."}, + {"source_table", std::make_shared(), "Name of the source table."}, + {"destination_database", std::make_shared(), "Name of the destination database."}, + {"destination_table", std::make_shared(), "Name of the destination table."}, + {"create_time", std::make_shared(), "Date and time when the export command was submitted."}, + {"partition_id", std::make_shared(), "ID of the partition being exported."}, + {"transaction_id", std::make_shared(), "ID of the export transaction (used to KILL it)."}, + {"query_id", std::make_shared(), "Query ID of the export operation."}, + {"parts", std::make_shared(std::make_shared()), "List of part names to be exported."}, + {"parts_count", std::make_shared(), "Number of parts in the export."}, + {"parts_to_do", std::make_shared(), "Number of parts still pending to be exported."}, + {"status", std::make_shared(), "Status of the export (PENDING, COMPLETED, FAILED, KILLED)."}, + {"last_exception", std::make_shared(), "Message of the most recent exception for this task, if any."}, + {"last_exception_part", std::make_shared(), "Part associated with the last exception (empty for task-level failures such as commit)."}, + {"last_exception_time", std::make_shared(), "Time of the most recent exception."}, + {"exception_count", std::make_shared(), "Total number of failures recorded for this task."}, + }; +} + +void StorageSystemPartitionExports::fillData(MutableColumns & res_columns, ContextPtr context, const ActionsDAG::Node * predicate, std::vector) const +{ + const auto access = context->getAccess(); + const bool check_access_for_databases = !access->isGranted(AccessType::SHOW_TABLES); + + std::map> merge_tree_tables; + for (const auto & db : DatabaseCatalog::instance().getDatabases(GetDatabasesOptions{.with_remote_databases = false})) + { + /// skip data lakes + if (db.second->isExternal()) + continue; + + const bool check_access_for_tables = check_access_for_databases && !access->isGranted(AccessType::SHOW_TABLES, db.first); + + for (auto iterator = db.second->getTablesIterator(context); iterator->isValid(); iterator->next()) + { + const auto & table = iterator->table(); + if (!table) + continue; + + /// Only plain (non-replicated) MergeTree tables expose partition exports here. + if (!dynamic_cast(table.get())) + continue; + + if (check_access_for_tables && !access->isGranted(AccessType::SHOW_TABLES, db.first, iterator->name())) + continue; + + merge_tree_tables[db.first][iterator->name()] = table; + } + } + + MutableColumnPtr col_database_mut = ColumnString::create(); + MutableColumnPtr col_table_mut = ColumnString::create(); + + for (auto & db : merge_tree_tables) + { + for (auto & table : db.second) + { + col_database_mut->insert(db.first); + col_table_mut->insert(table.first); + } + } + + ColumnPtr col_database = std::move(col_database_mut); + ColumnPtr col_table = std::move(col_table_mut); + + /// Determine what tables are needed by the conditions in the query. + { + Block filtered_block + { + { col_database, std::make_shared(), "database" }, + { col_table, std::make_shared(), "table" }, + }; + + VirtualColumnUtils::filterBlockWithPredicate(predicate, filtered_block, context); + + if (!filtered_block.rows()) + return; + + col_database = filtered_block.getByName("database").column; + col_table = filtered_block.getByName("table").column; + } + + for (size_t i_storage = 0; i_storage < col_database->size(); ++i_storage) + { + const auto database = (*col_database)[i_storage].safeGet(); + const auto table = (*col_table)[i_storage].safeGet(); + + std::vector partition_exports_info; + { + const IStorage * storage = merge_tree_tables[database][table].get(); + if (const auto * merge_tree = dynamic_cast(storage)) + partition_exports_info = merge_tree->getPartitionExportsInfo(); + } + + for (const PartitionExportInfo & info : partition_exports_info) + { + std::size_t i = 0; + res_columns[i++]->insert(info.source_database); + res_columns[i++]->insert(info.source_table); + res_columns[i++]->insert(info.destination_database); + res_columns[i++]->insert(info.destination_table); + res_columns[i++]->insert(info.create_time); + res_columns[i++]->insert(info.partition_id); + res_columns[i++]->insert(info.transaction_id); + res_columns[i++]->insert(info.query_id); + + Array parts_array; + parts_array.reserve(info.parts.size()); + for (const auto & part : info.parts) + parts_array.push_back(part); + res_columns[i++]->insert(parts_array); + + res_columns[i++]->insert(info.parts_count); + res_columns[i++]->insert(info.parts_to_do); + res_columns[i++]->insert(info.status); + res_columns[i++]->insert(info.last_exception_message); + res_columns[i++]->insert(info.last_exception_part); + res_columns[i++]->insert(info.last_exception_time); + res_columns[i++]->insert(info.exception_count); + } + } +} + +} diff --git a/src/Storages/System/StorageSystemPartitionExports.h b/src/Storages/System/StorageSystemPartitionExports.h new file mode 100644 index 000000000000..509f39718f09 --- /dev/null +++ b/src/Storages/System/StorageSystemPartitionExports.h @@ -0,0 +1,27 @@ +#pragma once + +#include + +namespace DB +{ + +class Context; + +/// system.partition_exports: progress of EXPORT PARTITION tasks on plain (non-replicated) MergeTree +/// tables. Backed by each table's on-disk task descriptors mirrored in memory; querying it does not +/// touch disk. Each export task is represented by a single row. (ReplicatedMergeTree tasks live in +/// system.replicated_partition_exports instead.) +class StorageSystemPartitionExports final : public IStorageSystemOneBlock +{ +public: + std::string getName() const override { return "SystemPartitionExports"; } + + static ColumnsDescription getColumnsDescription(); + +protected: + using IStorageSystemOneBlock::IStorageSystemOneBlock; + + void fillData(MutableColumns & res_columns, ContextPtr context, const ActionsDAG::Node *, std::vector) const override; +}; + +} diff --git a/src/Storages/System/attachSystemTables.cpp b/src/Storages/System/attachSystemTables.cpp index 57b6e4d3d7db..4e2511c9624f 100644 --- a/src/Storages/System/attachSystemTables.cpp +++ b/src/Storages/System/attachSystemTables.cpp @@ -1,5 +1,6 @@ #include #include +#include #include "config.h" #include @@ -243,6 +244,7 @@ void attachSystemTablesServer(ContextPtr context, IDatabase & system_database, b if (context->getServerSettings()[ServerSetting::allow_experimental_export_merge_tree_partition]) { attach(context, system_database, "replicated_partition_exports", "Contains a list of partition exports of ReplicatedMergeTree tables and their progress. Each export operation is represented by a single row."); + attach(context, system_database, "partition_exports", "Contains a list of partition exports of plain (non-replicated) MergeTree tables and their progress. Each export operation is represented by a single row."); } attach(context, system_database, "mutations", "Contains a list of mutations and their progress. Each mutation command is represented by a single row."); attachNoDescription(context, system_database, "replicas", "Contains information and status of all table replicas on current server. Each replica is represented by a single row."); diff --git a/tests/integration/helpers/export_partition_helpers.py b/tests/integration/helpers/export_partition_helpers.py index 46e73e8a04e6..123ea233b707 100644 --- a/tests/integration/helpers/export_partition_helpers.py +++ b/tests/integration/helpers/export_partition_helpers.py @@ -21,8 +21,12 @@ def wait_for_export_status( expected_status="COMPLETED", timeout=60, poll_interval=0.5, + system_table="replicated_partition_exports", ): - """Poll system.replicated_partition_exports until status matches. + """Poll a partition-exports system table until status matches. + + *system_table* selects the source table: ``replicated_partition_exports`` for + ReplicatedMergeTree (default) or ``partition_exports`` for plain MergeTree. *dest_table* may be ``None`` to skip filtering by destination table (useful for catalog-based tests where the destination is a database-qualified path). @@ -34,7 +38,7 @@ def wait_for_export_status( f" AND destination_table = '{dest_table}'" if dest_table else "" ) status = node.query( - f"SELECT status FROM system.replicated_partition_exports" + f"SELECT status FROM system.{system_table}" f" WHERE source_table = '{source_table}'" f"{dest_filter}" f" AND partition_id = '{partition_id}'" @@ -59,12 +63,13 @@ def wait_for_export_to_start( partition_id, timeout=10, poll_interval=0.2, + system_table="replicated_partition_exports", ): - """Poll until at least one row exists in system.replicated_partition_exports.""" + """Poll until at least one row exists in the given partition-exports system table.""" start_time = time.time() while time.time() - start_time < timeout: count = node.query( - f"SELECT count() FROM system.replicated_partition_exports" + f"SELECT count() FROM system.{system_table}" f" WHERE source_table = '{source_table}'" f" AND destination_table = '{dest_table}'" f" AND partition_id = '{partition_id}'" @@ -88,6 +93,7 @@ def wait_for_exception_count( min_exception_count=1, timeout=60, poll_interval=0.5, + system_table="replicated_partition_exports", ): """Wait for exception_count to reach at least *min_exception_count*. @@ -104,7 +110,7 @@ def wait_for_exception_count( last_exception_count = None while time.time() - start_time < timeout: exception_count_str = node.query( - f"SELECT exception_count FROM system.replicated_partition_exports" + f"SELECT exception_count FROM system.{system_table}" f" WHERE source_table = '{source_table}'" f" AND destination_table = '{dest_table}'" f" AND partition_id = '{partition_id}'" diff --git a/tests/integration/test_export_mt_partition_to_iceberg/__init__.py b/tests/integration/test_export_mt_partition_to_iceberg/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/integration/test_export_mt_partition_to_iceberg/configs/allow_experimental_export_partition.xml b/tests/integration/test_export_mt_partition_to_iceberg/configs/allow_experimental_export_partition.xml new file mode 100644 index 000000000000..514cd710836a --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_iceberg/configs/allow_experimental_export_partition.xml @@ -0,0 +1,3 @@ + + 1 + diff --git a/tests/integration/test_export_mt_partition_to_iceberg/configs/config.d/metadata_log.xml b/tests/integration/test_export_mt_partition_to_iceberg/configs/config.d/metadata_log.xml new file mode 100644 index 000000000000..c1fece21745c --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_iceberg/configs/config.d/metadata_log.xml @@ -0,0 +1,7 @@ + + + system + iceberg_metadata_log
+ 10 + + diff --git a/tests/integration/test_export_mt_partition_to_iceberg/configs/users.d/profile.xml b/tests/integration/test_export_mt_partition_to_iceberg/configs/users.d/profile.xml new file mode 100644 index 000000000000..8f9cbd70675d --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_iceberg/configs/users.d/profile.xml @@ -0,0 +1,7 @@ + + + + 3 + + + diff --git a/tests/integration/test_export_mt_partition_to_iceberg/test.py b/tests/integration/test_export_mt_partition_to_iceberg/test.py new file mode 100644 index 000000000000..2d67d99d8b22 --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_iceberg/test.py @@ -0,0 +1,114 @@ +import logging + +import pytest + +from helpers.cluster import ClickHouseCluster +from helpers.export_partition_helpers import ( + make_iceberg_s3, + make_mt, + unique_suffix, + wait_for_export_status, +) + +# Plain (non-replicated) MergeTree EXPORT PARTITION to an Iceberg (IcebergS3, catalog-less) +# destination. The Iceberg commit path is shared with the replicated implementation; these tests +# focus on the plain-MergeTree entry point and EXPORT PARTITION ALL. +# +# Restart-resume is deliberately NOT retested here: it is destination-agnostic (driven by the +# scheduler + on-disk descriptor) and is covered by +# test_export_mt_partition_to_object_storage::test_export_partition_resumes_after_restart. It cannot +# be reproduced the same way for Iceberg anyway, because an Iceberg destination reads its +# metadata.json from object storage at request time, so blocking object storage up-front would make +# the ALTER itself fail instead of leaving a resumable in-flight task. + +SYSTEM_TABLE = "partition_exports" + + +@pytest.fixture(scope="module") +def cluster(): + try: + cluster = ClickHouseCluster(__file__) + cluster.add_instance( + "node", + main_configs=[ + "configs/allow_experimental_export_partition.xml", + "configs/config.d/metadata_log.xml", + ], + user_configs=["configs/users.d/profile.xml"], + with_minio=True, + stay_alive=True, + ) + logging.info("Starting cluster...") + cluster.start() + yield cluster + finally: + cluster.shutdown() + + +@pytest.fixture(autouse=True) +def drop_tables_after_test(cluster): + yield + for instance_name, instance in cluster.instances.items(): + try: + tables_str = instance.query( + "SELECT name FROM system.tables WHERE database = 'default' FORMAT TabSeparated" + ).strip() + if not tables_str: + continue + for table in tables_str.split("\n"): + table = table.strip() + if table: + instance.query(f"DROP TABLE IF EXISTS default.`{table}` SYNC") + except Exception as e: + logging.warning(f"drop_tables_after_test: cleanup failed on {instance_name}: {e}") + + +def setup_tables(node, mt_table, iceberg_table): + make_mt(node, mt_table, "id Int64, year Int32", partition_by="year") + node.query(f"INSERT INTO {mt_table} VALUES (1, 2020), (2, 2020), (3, 2020), (4, 2021)") + make_iceberg_s3(node, iceberg_table, "id Int64, year Int32", partition_by="year") + + +def test_export_partition_to_iceberg(cluster): + node = cluster.instances["node"] + + uid = unique_suffix() + mt_table = f"mt_{uid}" + iceberg_table = f"iceberg_{uid}" + + setup_tables(node, mt_table, iceberg_table) + + node.query( + f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {iceberg_table}", + settings={"allow_insert_into_iceberg": 1}, + ) + wait_for_export_status(node, mt_table, iceberg_table, "2020", "COMPLETED", system_table=SYSTEM_TABLE) + + count = int(node.query(f"SELECT count() FROM {iceberg_table}").strip()) + assert count == 3, f"Expected 3 rows in Iceberg table after export, got {count}" + + result = node.query(f"SELECT id, year FROM {iceberg_table} ORDER BY id").strip() + assert result == "1\t2020\n2\t2020\n3\t2020", f"Unexpected data in Iceberg table:\n{result}" + + +def test_export_partition_all_to_iceberg(cluster): + node = cluster.instances["node"] + + uid = unique_suffix() + mt_table = f"mt_all_{uid}" + iceberg_table = f"iceberg_all_{uid}" + + setup_tables(node, mt_table, iceberg_table) + + node.query( + f"ALTER TABLE {mt_table} EXPORT PARTITION ALL TO TABLE {iceberg_table}", + settings={"allow_insert_into_iceberg": 1}, + ) + + wait_for_export_status(node, mt_table, iceberg_table, "2020", "COMPLETED", system_table=SYSTEM_TABLE) + wait_for_export_status(node, mt_table, iceberg_table, "2021", "COMPLETED", system_table=SYSTEM_TABLE) + + count_2020 = int(node.query(f"SELECT count() FROM {iceberg_table} WHERE year = 2020").strip()) + count_2021 = int(node.query(f"SELECT count() FROM {iceberg_table} WHERE year = 2021").strip()) + assert count_2020 == 3, f"Expected 3 rows for year=2020, got {count_2020}" + assert count_2021 == 1, f"Expected 1 row for year=2021, got {count_2021}" diff --git a/tests/integration/test_export_mt_partition_to_object_storage/__init__.py b/tests/integration/test_export_mt_partition_to_object_storage/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/integration/test_export_mt_partition_to_object_storage/configs/allow_experimental_export_partition.xml b/tests/integration/test_export_mt_partition_to_object_storage/configs/allow_experimental_export_partition.xml new file mode 100644 index 000000000000..514cd710836a --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_object_storage/configs/allow_experimental_export_partition.xml @@ -0,0 +1,3 @@ + + 1 + diff --git a/tests/integration/test_export_mt_partition_to_object_storage/configs/disable_experimental_export_partition.xml b/tests/integration/test_export_mt_partition_to_object_storage/configs/disable_experimental_export_partition.xml new file mode 100644 index 000000000000..21be4dc894df --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_object_storage/configs/disable_experimental_export_partition.xml @@ -0,0 +1,3 @@ + + 0 + diff --git a/tests/integration/test_export_mt_partition_to_object_storage/configs/named_collections.xml b/tests/integration/test_export_mt_partition_to_object_storage/configs/named_collections.xml new file mode 100644 index 000000000000..573822539c50 --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_object_storage/configs/named_collections.xml @@ -0,0 +1,9 @@ + + + + http://minio1:9001/root/data + minio + ClickHouse_Minio_P@ssw0rd + + + diff --git a/tests/integration/test_export_mt_partition_to_object_storage/configs/users.d/profile.xml b/tests/integration/test_export_mt_partition_to_object_storage/configs/users.d/profile.xml new file mode 100644 index 000000000000..8f9cbd70675d --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_object_storage/configs/users.d/profile.xml @@ -0,0 +1,7 @@ + + + + 3 + + + diff --git a/tests/integration/test_export_mt_partition_to_object_storage/test.py b/tests/integration/test_export_mt_partition_to_object_storage/test.py new file mode 100644 index 000000000000..717bac90dc7a --- /dev/null +++ b/tests/integration/test_export_mt_partition_to_object_storage/test.py @@ -0,0 +1,307 @@ +import logging +import time +import uuid + +import pytest + +from helpers.cluster import ClickHouseCluster +from helpers.export_partition_helpers import ( + unique_suffix, + wait_for_export_status, + wait_for_export_to_start, +) +from helpers.network import PartitionManager + +# Plain (non-replicated) MergeTree EXPORT PARTITION. Unlike the replicated variant there is no +# ZooKeeper coordination: the task descriptor is persisted on the table's disk and a local +# background scheduler drives it, so these tests also cover restart-resume. + +SYSTEM_TABLE = "partition_exports" + + +def skip_if_remote_database_disk_enabled(cluster): + for instance in cluster.instances.values(): + if instance.with_remote_database_disk: + pytest.skip( + "Test cannot run with remote database disk enabled, as it blocks MinIO which stores database metadata" + ) + + +@pytest.fixture(scope="module") +def cluster(): + try: + cluster = ClickHouseCluster(__file__) + cluster.add_instance( + "node", + main_configs=[ + "configs/named_collections.xml", + "configs/allow_experimental_export_partition.xml", + ], + user_configs=["configs/users.d/profile.xml"], + with_minio=True, + stay_alive=True, + ) + cluster.add_instance( + "node_export_disabled", + main_configs=[ + "configs/named_collections.xml", + "configs/disable_experimental_export_partition.xml", + ], + user_configs=["configs/users.d/profile.xml"], + with_minio=True, + stay_alive=True, + ) + logging.info("Starting cluster...") + cluster.start() + yield cluster + finally: + cluster.shutdown() + + +@pytest.fixture(autouse=True) +def drop_tables_after_test(cluster): + yield + for instance_name, instance in cluster.instances.items(): + try: + tables_str = instance.query( + "SELECT name FROM system.tables WHERE database = 'default' FORMAT TabSeparated" + ).strip() + if not tables_str: + continue + for table in tables_str.split("\n"): + table = table.strip() + if table: + instance.query(f"DROP TABLE IF EXISTS default.`{table}` SYNC") + except Exception as e: + logging.warning(f"drop_tables_after_test: cleanup failed on {instance_name}: {e}") + + +def create_s3_table(node, s3_table): + node.query( + f"CREATE TABLE {s3_table} (id UInt64, year UInt16) " + f"ENGINE = S3(s3_conn, filename='{s3_table}', format=Parquet, partition_strategy='hive') " + f"PARTITION BY year" + ) + + +def create_tables_and_insert_data(node, mt_table, s3_table): + node.query(f"DROP TABLE IF EXISTS {mt_table} SYNC") + node.query( + f"CREATE TABLE {mt_table} (id UInt64, year UInt16) ENGINE = MergeTree " + f"PARTITION BY year ORDER BY tuple() " + f"SETTINGS enable_block_number_column = 1, enable_block_offset_column = 1" + ) + node.query(f"INSERT INTO {mt_table} VALUES (1, 2020), (2, 2020), (3, 2020), (4, 2021)") + create_s3_table(node, s3_table) + + +def test_export_partition_basic_roundtrip(cluster): + node = cluster.instances["node"] + + postfix = unique_suffix() + mt_table = f"basic_mt_{postfix}" + s3_table = f"basic_s3_{postfix}" + + create_tables_and_insert_data(node, mt_table, s3_table) + + node.query(f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {s3_table}") + + wait_for_export_status(node, mt_table, s3_table, "2020", "COMPLETED", system_table=SYSTEM_TABLE) + + assert node.query(f"SELECT count() FROM {s3_table} WHERE year = 2020") == "3\n" + assert ( + node.query( + f"SELECT count() FROM s3(s3_conn, filename='{s3_table}/commit_2020_*', format=LineAsString)" + ) + != "0\n" + ), "Commit file missing for partition 2020" + + +def test_export_partition_all(cluster): + node = cluster.instances["node"] + + postfix = unique_suffix() + mt_table = f"all_mt_{postfix}" + s3_table = f"all_s3_{postfix}" + + node.query( + f"CREATE TABLE {mt_table} (id UInt64, year UInt16) ENGINE = MergeTree " + f"PARTITION BY year ORDER BY tuple()" + ) + node.query(f"INSERT INTO {mt_table} VALUES (1, 2020), (2, 2021), (3, 2022)") + create_s3_table(node, s3_table) + + node.query(f"ALTER TABLE {mt_table} EXPORT PARTITION ALL TO TABLE {s3_table}") + + for partition_id in ("2020", "2021", "2022"): + wait_for_export_status(node, mt_table, s3_table, partition_id, "COMPLETED", system_table=SYSTEM_TABLE) + + assert node.query(f"SELECT count() FROM {s3_table}") == "3\n" + + +def test_duplicate_export_rejected_and_force(cluster): + node = cluster.instances["node"] + + postfix = unique_suffix() + mt_table = f"dup_mt_{postfix}" + s3_table = f"dup_s3_{postfix}" + + create_tables_and_insert_data(node, mt_table, s3_table) + + node.query(f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {s3_table}") + wait_for_export_status(node, mt_table, s3_table, "2020", "COMPLETED", system_table=SYSTEM_TABLE) + + # Re-export without force must be rejected. + error = node.query_and_get_error( + f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {s3_table}" + ) + assert "EXPORT_PARTITION_ALREADY_EXPORTED" in error, f"Expected duplicate rejection, got: {error}" + + # With force + overwrite policy the re-export succeeds. + node.query( + f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {s3_table} " + f"SETTINGS export_merge_tree_partition_force_export = 1, " + f"export_merge_tree_part_file_already_exists_policy = 'overwrite'" + ) + wait_for_export_status(node, mt_table, s3_table, "2020", "COMPLETED", system_table=SYSTEM_TABLE) + + +def test_kill_export(cluster): + skip_if_remote_database_disk_enabled(cluster) + node = cluster.instances["node"] + + postfix = unique_suffix() + mt_table = f"kill_mt_{postfix}" + s3_table = f"kill_s3_{postfix}" + + create_tables_and_insert_data(node, mt_table, s3_table) + + minio_ip = cluster.minio_ip + minio_port = cluster.minio_port + + with PartitionManager() as pm: + pm.add_rule({ + "instance": node, + "destination": node.ip_address, + "protocol": "tcp", + "source_port": minio_port, + "action": "REJECT --reject-with tcp-reset", + }) + pm.add_rule({ + "instance": node, + "destination": minio_ip, + "protocol": "tcp", + "destination_port": minio_port, + "action": "REJECT --reject-with tcp-reset", + }) + + node.query(f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {s3_table}") + wait_for_export_to_start(node, mt_table, s3_table, "2020", system_table=SYSTEM_TABLE) + + node.query( + f"KILL EXPORT PARTITION WHERE partition_id = '2020'" + f" AND source_table = '{mt_table}' AND destination_table = '{s3_table}'" + ) + + wait_for_export_status(node, mt_table, s3_table, "2020", "KILLED", system_table=SYSTEM_TABLE, timeout=30) + + # No commit file and no data must have landed. + assert ( + node.query( + f"SELECT count() FROM s3(s3_conn, filename='{s3_table}/commit_2020_*', format=LineAsString)" + ) + == "0\n" + ), "Partition 2020 was committed despite being killed" + assert node.query(f"SELECT count() FROM {s3_table} WHERE year = 2020") == "0\n" + + +def test_export_partition_resumes_after_restart(cluster): + """The distinguishing feature of the plain MergeTree implementation: the on-disk task + descriptor must let an in-flight export resume after a hard restart.""" + skip_if_remote_database_disk_enabled(cluster) + node = cluster.instances["node"] + + postfix = unique_suffix() + mt_table = f"restart_mt_{postfix}" + s3_table = f"restart_s3_{postfix}" + + create_tables_and_insert_data(node, mt_table, s3_table) + + minio_ip = cluster.minio_ip + minio_port = cluster.minio_port + + with PartitionManager() as pm: + pm.add_rule({ + "instance": node, + "destination": node.ip_address, + "protocol": "tcp", + "source_port": minio_port, + "action": "REJECT --reject-with tcp-reset", + }) + pm.add_rule({ + "instance": node, + "destination": minio_ip, + "protocol": "tcp", + "destination_port": minio_port, + "action": "REJECT --reject-with tcp-reset", + }) + + node.query(f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {s3_table}") + wait_for_export_to_start(node, mt_table, s3_table, "2020", system_table=SYSTEM_TABLE) + + # Kill the server while the export is still in flight (S3 blocked, nothing committed yet). + node.stop_clickhouse(kill=True) + + # We cannot observe the "nothing committed before the crash" invariant on a single node: while + # the only node is down there is nothing to query S3 with, and once it restarts (S3 now + # reachable) the persisted PENDING task resumes immediately. So we only assert the actual + # restart-resume behavior: the task must resume from its on-disk descriptor and complete. + node.start_clickhouse() + + wait_for_export_status(node, mt_table, s3_table, "2020", "COMPLETED", system_table=SYSTEM_TABLE, timeout=90) + assert node.query(f"SELECT count() FROM {s3_table} WHERE year = 2020") == "3\n" + + +def test_export_partition_feature_disabled(cluster): + node = cluster.instances["node_export_disabled"] + + postfix = unique_suffix() + mt_table = f"disabled_mt_{postfix}" + s3_table = f"disabled_s3_{postfix}" + + create_tables_and_insert_data(node, mt_table, s3_table) + + error = node.query_and_get_error( + f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {s3_table}" + ) + assert "experimental" in error, f"Expected experimental-feature error, got: {error}" + + error = node.query_and_get_error( + f"KILL EXPORT PARTITION WHERE partition_id = '2020' AND source_table = '{mt_table}'" + f" AND destination_table = '{s3_table}'" + ) + assert "experimental" in error, f"Expected experimental-feature error on KILL, got: {error}" + + +def test_pending_mutations_throw_before_export(cluster): + node = cluster.instances["node"] + + postfix = unique_suffix() + mt_table = f"pending_mut_mt_{postfix}" + s3_table = f"pending_mut_s3_{postfix}" + + create_tables_and_insert_data(node, mt_table, s3_table) + + node.query(f"SYSTEM STOP MERGES {mt_table}") + node.query(f"ALTER TABLE {mt_table} UPDATE id = id + 100 WHERE year = 2020") + + mutations = node.query( + f"SELECT count() FROM system.mutations WHERE table = '{mt_table}' AND is_done = 0" + ) + assert mutations.strip() != "0", "Mutation should be pending" + + error = node.query_and_get_error( + f"ALTER TABLE {mt_table} EXPORT PARTITION ID '2020' TO TABLE {s3_table} " + f"SETTINGS export_merge_tree_part_throw_on_pending_mutations = true" + ) + assert "PENDING_MUTATIONS_NOT_ALLOWED" in error, f"Expected pending-mutations error, got: {error}" diff --git a/tests/queries/0_stateless/04303_export_plain_merge_tree_partition.reference b/tests/queries/0_stateless/04303_export_plain_merge_tree_partition.reference new file mode 100644 index 000000000000..88ff7663c0fe --- /dev/null +++ b/tests/queries/0_stateless/04303_export_plain_merge_tree_partition.reference @@ -0,0 +1,31 @@ +Export partition 2020 +Export partition 2021 +Select from destination table (2020, 2021) +1 2020 +2 2020 +3 2020 +4 2021 +5 2021 +Re-exporting 2020 without force is rejected +EXPORT_PARTITION_ALREADY_EXPORTED +Export remaining partitions with EXPORT PARTITION ALL (skip existing) +Select from destination table (all partitions) +1 2020 +2 2020 +3 2020 +4 2021 +5 2021 +6 2022 +7 2022 +Roundtrip: create a MergeTree table from the exported S3 data +1 2020 +2 2020 +3 2020 +4 2021 +5 2021 +6 2022 +7 2022 +system.partition_exports statuses +2020 COMPLETED 2 0 +2021 COMPLETED 2 0 +2022 COMPLETED 1 0 diff --git a/tests/queries/0_stateless/04303_export_plain_merge_tree_partition.sh b/tests/queries/0_stateless/04303_export_plain_merge_tree_partition.sh new file mode 100644 index 000000000000..a744dc5debd0 --- /dev/null +++ b/tests/queries/0_stateless/04303_export_plain_merge_tree_partition.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# Tags: no-fasttest, no-shared-merge-tree +# no-fasttest: requires S3 / MinIO. +# no-shared-merge-tree: this test exercises EXPORT PARTITION on a plain (non-replicated) MergeTree. + +CUR_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +# shellcheck source=../shell_config.sh +. "$CUR_DIR"/../shell_config.sh + +mt_table="mt_table_${CLICKHOUSE_DATABASE}" +s3_table="s3_table_${CLICKHOUSE_DATABASE}" +mt_roundtrip="mt_roundtrip_${CLICKHOUSE_DATABASE}" + +query() { + $CLICKHOUSE_CLIENT --query "$1" +} + +# Poll system.partition_exports until the given partition reaches the expected status (or timeout). +wait_for_status() { + local partition_id="$1" + local expected="$2" + local i=0 + while [ "$i" -lt 120 ]; do + status=$(query "SELECT status FROM system.partition_exports WHERE source_table = '$mt_table' AND destination_table = '$s3_table' AND partition_id = '$partition_id'") + if [ "$status" = "$expected" ]; then + return 0 + fi + sleep 0.5 + i=$((i + 1)) + done + echo "TIMEOUT waiting for partition $partition_id to reach $expected (last: '$status')" + return 1 +} + +query "DROP TABLE IF EXISTS $mt_table" +query "DROP TABLE IF EXISTS $s3_table" +query "DROP TABLE IF EXISTS $mt_roundtrip" + +query "CREATE TABLE $mt_table (id UInt64, year UInt16) ENGINE = MergeTree PARTITION BY year ORDER BY tuple()" +query "CREATE TABLE $s3_table (id UInt64, year UInt16) ENGINE = S3(s3_conn, filename='$s3_table', format=Parquet, partition_strategy='hive') PARTITION BY year" + +# Stop merges so the number of parts per partition stays stable for the assertions below. +query "SYSTEM STOP MERGES $mt_table" + +query "INSERT INTO $mt_table VALUES (1, 2020), (2, 2020), (4, 2021)" +query "INSERT INTO $mt_table VALUES (3, 2020), (5, 2021)" +query "INSERT INTO $mt_table VALUES (6, 2022), (7, 2022)" + +echo "Export partition 2020" +query "ALTER TABLE $mt_table EXPORT PARTITION ID '2020' TO TABLE $s3_table" +wait_for_status "2020" "COMPLETED" + +echo "Export partition 2021" +query "ALTER TABLE $mt_table EXPORT PARTITION ID '2021' TO TABLE $s3_table" +wait_for_status "2021" "COMPLETED" + +echo "Select from destination table (2020, 2021)" +query "SELECT * FROM $s3_table ORDER BY id" + +echo "Re-exporting 2020 without force is rejected" +query "ALTER TABLE $mt_table EXPORT PARTITION ID '2020' TO TABLE $s3_table" 2>&1 | grep -o "EXPORT_PARTITION_ALREADY_EXPORTED" | head -1 + +echo "Export remaining partitions with EXPORT PARTITION ALL (skip existing)" +query "ALTER TABLE $mt_table EXPORT PARTITION ALL TO TABLE $s3_table SETTINGS export_merge_tree_partition_all_on_error = 'skip_conflicts'" +wait_for_status "2022" "COMPLETED" + +echo "Select from destination table (all partitions)" +query "SELECT * FROM $s3_table ORDER BY id" + +echo "Roundtrip: create a MergeTree table from the exported S3 data" +query "CREATE TABLE $mt_roundtrip ENGINE = MergeTree PARTITION BY year ORDER BY tuple() AS SELECT * FROM $s3_table" +query "SELECT * FROM $mt_roundtrip ORDER BY id" + +echo "system.partition_exports statuses" +query "SELECT partition_id, status, parts_count, parts_to_do FROM system.partition_exports WHERE source_table = '$mt_table' AND destination_table = '$s3_table' ORDER BY partition_id" + +query "DROP TABLE IF EXISTS $mt_table" +query "DROP TABLE IF EXISTS $s3_table" +query "DROP TABLE IF EXISTS $mt_roundtrip"