almost done with kill export partition

Author: arthurpassos

src/Interpreters/InterpreterKillQueryQuery.cpp

@@ -250,6 +250,77 @@ BlockIO InterpreterKillQueryQuery::execute()
 
         break;
     }
+    case ASTKillQueryQuery::Type::ExportPartition:
+    {
+        Block exports_block = getSelectResult(
+            "source_database, source_table, transaction_id, destination_database, destination_table, partition_id",
+            "system.replicated_partition_exports");
+        if (!exports_block)
+            return res_io;
+
+        const ColumnString & src_db_col = typeid_cast<const ColumnString &>(*exports_block.getByName("source_database").column);
+        const ColumnString & src_tbl_col = typeid_cast<const ColumnString &>(*exports_block.getByName("source_table").column);
+        const ColumnString & dst_db_col = typeid_cast<const ColumnString &>(*exports_block.getByName("destination_database").column);
+        const ColumnString & dst_tbl_col = typeid_cast<const ColumnString &>(*exports_block.getByName("destination_table").column);
+        const ColumnString & tx_col = typeid_cast<const ColumnString &>(*exports_block.getByName("transaction_id").column);
+
+        auto header = exports_block.cloneEmpty();
+        header.insert(0, {ColumnString::create(), std::make_shared<DataTypeString>(), "kill_status"});
+
+        MutableColumns res_columns = header.cloneEmptyColumns();
+        auto table_id = StorageID::createEmpty();
+        AccessRightsElements required_access_rights;
+        auto access = getContext()->getAccess();
+        bool access_denied = false;
+
+        for (size_t i = 0; i < exports_block.rows(); ++i)
+        {
+            const auto src_database = src_db_col.getDataAt(i).toString();
+            const auto src_table = src_tbl_col.getDataAt(i).toString();
+            const auto dst_database = dst_db_col.getDataAt(i).toString();
+            const auto dst_table = dst_tbl_col.getDataAt(i).toString();
+
+            table_id = StorageID{src_database, src_table};
+            auto transaction_id = tx_col.getDataAt(i).toString();
+
+            CancellationCode code = CancellationCode::Unknown;
+            if (!query.test)
+            {
+                auto storage = DatabaseCatalog::instance().tryGetTable(table_id, getContext());
+                if (!storage)
+                    code = CancellationCode::NotFound;
+                else
+                {
+                    ASTAlterCommand alter_command{};
+                    alter_command.type = ASTAlterCommand::EXPORT_PARTITION;
+                    alter_command.move_destination_type = DataDestinationType::TABLE;
+                    alter_command.from_database = src_database;
+                    alter_command.from_table = src_table;
+                    alter_command.to_database = dst_database;
+                    alter_command.to_table = dst_table;
+
+                    required_access_rights = InterpreterAlterQuery::getRequiredAccessForCommand(
+                        alter_command, table_id.database_name, table_id.table_name);
+                    if (!access->isGranted(required_access_rights))
+                    {
+                        access_denied = true;
+                        continue;
+                    }
+                    code = storage->killExportPartition(transaction_id);
+                }
+            }
+
+            insertResultRow(i, code, exports_block, header, res_columns);
+        }
+
+        if (res_columns[0]->empty() && access_denied)
+            throw Exception(ErrorCodes::ACCESS_DENIED, "Not allowed to kill export partition. "
+                "To execute this query, it's necessary to have the grant {}", required_access_rights.toString());
+
+        res_io.pipeline = QueryPipeline(Pipe(std::make_shared<SourceFromSingleChunk>(header.cloneWithColumns(std::move(res_columns)))));
+
+        break;
+    }
     case ASTKillQueryQuery::Type::Mutation:
     {
         Block mutations_block = getSelectResult("database, table, mutation_id, command", "system.mutations");
@@ -462,6 +533,9 @@ AccessRightsElements InterpreterKillQueryQuery::getRequiredAccessForDDLOnCluster
                 | AccessType::ALTER_MATERIALIZE_COLUMN
                 | AccessType::ALTER_MATERIALIZE_TTL
             );
+    /// todo arthur think about this
+    else if (query.type == ASTKillQueryQuery::Type::ExportPartition)
+        required_access.emplace_back(AccessType::ALTER_EXPORT_PARTITION);
     return required_access;
 }
 

src/Parsers/ASTKillQueryQuery.cpp

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

src/Parsers/ASTKillQueryQuery.h

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

src/Parsers/ParserKillQueryQuery.cpp

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

src/Storages/IStorage.cpp

@@ -300,6 +300,11 @@ CancellationCode IStorage::killPartMoveToShard(const UUID & /*task_uuid*/)
     throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Part moves between shards are not supported by storage {}", getName());
 }
 
+CancellationCode IStorage::killExportPartition(const String & /*transaction_id*/)
+{
+    throw Exception(ErrorCodes::NOT_IMPLEMENTED, "Export partition is not supported by storage {}", getName());
+}
+
 StorageID IStorage::getStorageID() const
 {
     std::lock_guard lock(id_mutex);

src/Storages/IStorage.h

@@ -574,6 +574,9 @@ It is currently only implemented in StorageObjectStorage.
 
     virtual void setMutationCSN(const String & /*mutation_id*/, UInt64 /*csn*/);
 
+    /// Cancel a replicated partition export by transaction id.
+    virtual CancellationCode killExportPartition(const String & /*transaction_id*/);
+
     /// Cancel a part move to shard.
     virtual CancellationCode killPartMoveToShard(const UUID & /*task_uuid*/);
 

src/Storages/MergeTree/MergeTreeData.cpp

@@ -5910,14 +5910,15 @@ void MergeTreeData::exportPartToTable(const PartitionCommand & command, ContextP
 
     const auto database_name = query_context->resolveDatabase(command.to_database);
 
-    exportPartToTable(part_name, StorageID{database_name, command.to_table}, query_context);
+    exportPartToTable(part_name, StorageID{database_name, command.to_table}, query_context->getCurrentQueryId(), query_context);
 }
 
 void MergeTreeData::exportPartToTable(
     const std::string & part_name,
     const StorageID & destination_storage_id,
+    const String & transaction_id,
     ContextPtr query_context,
-    std::function<void(MergeTreeExportManifest::CompletionCallbackResult)> completion_callback)
+    std::function<void(MergeTreePartExportManifest::CompletionCallbackResult)> completion_callback)
 {
     auto dest_storage = DatabaseCatalog::instance().getTable(destination_storage_id, query_context);
 
@@ -5950,9 +5951,10 @@ void MergeTreeData::exportPartToTable(
                         part_name, getStorageID().getFullTableName());
 
     {
-        MergeTreeExportManifest manifest(
+        MergeTreePartExportManifest manifest(
             dest_storage->getStorageID(),
             part,
+            transaction_id,
             query_context->getSettingsRef()[Setting::export_merge_tree_part_overwrite_file_if_exists],
             query_context->getSettingsRef()[Setting::output_format_parallel_formatting],
             completion_callback);
@@ -5970,7 +5972,7 @@ void MergeTreeData::exportPartToTable(
 }
 
 void MergeTreeData::exportPartToTableImpl(
-    const MergeTreeExportManifest & manifest,
+    const MergeTreePartExportManifest & manifest,
     ContextPtr local_context)
 {
     auto metadata_snapshot = getInMemoryMetadataPtr();
@@ -6027,7 +6029,7 @@ void MergeTreeData::exportPartToTableImpl(
         export_manifests.erase(manifest);
 
         if (manifest.completion_callback)
-            manifest.completion_callback(MergeTreeExportManifest::CompletionCallbackResult::createFailure(e.message()));
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createFailure(e.message()));
         return;
     }
 
@@ -6096,6 +6098,9 @@ void MergeTreeData::exportPartToTableImpl(
 
     pipeline.complete(sink);
 
+    /// oh boy, is there another way?
+    manifest.pipeline = &pipeline;
+
     try
     {
         CompletedPipelineExecutor exec(pipeline);
@@ -6119,7 +6124,7 @@ void MergeTreeData::exportPartToTableImpl(
         ProfileEvents::increment(ProfileEvents::PartsExportTotalMilliseconds, static_cast<UInt64>((*exports_list_entry)->elapsed * 1000));
 
         if (manifest.completion_callback)
-            manifest.completion_callback(MergeTreeExportManifest::CompletionCallbackResult::createSuccess(destination_file_path));
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createSuccess(destination_file_path));
     }
     catch (const Exception & e)
     {
@@ -6142,12 +6147,30 @@ void MergeTreeData::exportPartToTableImpl(
         export_manifests.erase(manifest);
 
         if (manifest.completion_callback)
-            manifest.completion_callback(MergeTreeExportManifest::CompletionCallbackResult::createFailure(e.message()));
+            manifest.completion_callback(MergeTreePartExportManifest::CompletionCallbackResult::createFailure(e.message()));
 
         throw;
     }
 }
 
+void MergeTreeData::killExportPart(const String & query_id)
+{
+    std::lock_guard lock(export_manifests_mutex);
+
+    const auto it = std::find_if(export_manifests.begin(), export_manifests.end(), [&](const auto & manifest)
+    {
+        return manifest.query_id == query_id;
+    });
+
+    if (it == export_manifests.end())
+        return;
+
+    if (it->pipeline)
+        it->pipeline->cancel();
+
+    export_manifests.erase(it);
+}
+
 void MergeTreeData::movePartitionToShard(const ASTPtr & /*partition*/, bool /*move_part*/, const String & /*to*/, ContextPtr /*query_context*/)
 {
     throw Exception(ErrorCodes::NOT_IMPLEMENTED, "MOVE PARTITION TO SHARD is not supported by storage {}", getName());
@@ -8881,7 +8904,9 @@ bool MergeTreeData::scheduleDataMovingJob(BackgroundJobsAssignee & assignee)
         }
 
         manifest.in_progress = assignee.scheduleMoveTask(std::make_shared<ExecutableLambdaAdapter>(
-            [this, manifest] () mutable {
+            [this, &manifest] () mutable {
+                /// TODO arthur fix this: I need to be able to modify the real manifest
+                /// but grabbing it by reference is causing problems
                 exportPartToTableImpl(manifest, getContext());
                 return true;
             },

src/Storages/MergeTree/MergeTreeData.h

@@ -35,8 +35,8 @@
 #include <Interpreters/PartLog.h>
 #include <Poco/Timestamp.h>
 #include <Common/threadPoolCallbackRunner.h>
-#include <Storages/MergeTree/MergeTreeExportStatus.h>
-#include <Storages/MergeTree/MergeTreeExportManifest.h>
+#include <Storages/MergeTree/MergeTreePartExportStatus.h>
+#include <Storages/MergeTree/MergeTreePartExportManifest.h>
 
 #include <boost/multi_index_container.hpp>
 #include <boost/multi_index/ordered_index.hpp>
@@ -908,16 +908,20 @@ class MergeTreeData : public IStorage, public WithMutableContext
 
     void exportPartToTable(
         const std::string & part_name,
-        const StorageID & destination_storage_id, ContextPtr query_context,
-        std::function<void(MergeTreeExportManifest::CompletionCallbackResult)> completion_callback = {});
+        const StorageID & destination_storage_id,
+        const String & transaction_id,
+        ContextPtr query_context,
+        std::function<void(MergeTreePartExportManifest::CompletionCallbackResult)> completion_callback = {});
+
+    void killExportPart(const String & query_id);
 
     virtual void exportPartitionToTable(const PartitionCommand &, ContextPtr)
     {
         throw Exception(ErrorCodes::NOT_IMPLEMENTED, "EXPORT PARTITION is not implemented");
     }
 
     void exportPartToTableImpl(
-        const MergeTreeExportManifest & manifest,
+        const MergeTreePartExportManifest & manifest,
         ContextPtr local_context);
 
     /// Checks that Partition could be dropped right now
@@ -1166,7 +1170,7 @@ class MergeTreeData : public IStorage, public WithMutableContext
 
     mutable std::mutex export_manifests_mutex;
 
-    std::set<MergeTreeExportManifest> export_manifests;
+    std::set<MergeTreePartExportManifest> export_manifests;
 
     PinnedPartUUIDsPtr getPinnedPartUUIDs() const;
 

…ages/MergeTree/MergeTreeExportManifest.h …/MergeTree/MergeTreePartExportManifest.hsrc/Storages/MergeTree/MergeTreeExportManifest.h renamed to src/Storages/MergeTree/MergeTreePartExportManifest.h src/Storages/MergeTree/MergeTreeExportManifest.h renamed to src/Storages/MergeTree/MergeTreePartExportManifest.h

@@ -1,10 +1,11 @@
 #include <Interpreters/StorageID.h>
 #include <Storages/MergeTree/IMergeTreeDataPart.h>
+#include "QueryPipeline/QueryPipeline.h"
 
 namespace DB
 {
 
-struct MergeTreeExportManifest
+struct MergeTreePartExportManifest
 {
     using DataPartPtr = std::shared_ptr<const IMergeTreeDataPart>;
 
@@ -30,21 +31,25 @@ struct MergeTreeExportManifest
         String exception;
     };
 
-    MergeTreeExportManifest(
+    MergeTreePartExportManifest(
         const StorageID & destination_storage_id_,
         const DataPartPtr & data_part_,
+        const String & query_id_,
         bool overwrite_file_if_exists_,
         bool parallel_formatting_,
         std::function<void(CompletionCallbackResult)> completion_callback_ = {})
         : destination_storage_id(destination_storage_id_),
           data_part(data_part_),
+          query_id(query_id_),
           overwrite_file_if_exists(overwrite_file_if_exists_),
           parallel_formatting(parallel_formatting_),
           completion_callback(completion_callback_),
           create_time(time(nullptr)) {}
 
     StorageID destination_storage_id;
     DataPartPtr data_part;
+    /// Used for killing the export.
+    String query_id;
     bool overwrite_file_if_exists;
     bool parallel_formatting;
 
@@ -53,8 +58,10 @@ struct MergeTreeExportManifest
 
     time_t create_time;
     mutable bool in_progress = false;
+    /// Used for killing the export
+    mutable QueryPipeline * pipeline = nullptr;
 
-    bool operator<(const MergeTreeExportManifest & rhs) const 
+    bool operator<(const MergeTreePartExportManifest & rhs) const 
     {
         // Lexicographic comparison: first compare destination storage, then part name
         auto lhs_storage = destination_storage_id.getQualifiedName();
@@ -66,7 +73,7 @@ struct MergeTreeExportManifest
         return data_part->name < rhs.data_part->name;
     }
 
-    bool operator==(const MergeTreeExportManifest & rhs) const 
+    bool operator==(const MergeTreePartExportManifest & rhs) const 
     {
         return destination_storage_id.getQualifiedName() == rhs.destination_storage_id.getQualifiedName()
             && data_part->name == rhs.data_part->name;