CLICKHOUSE-3714 add max_partition_size_to_drop

dbms/programs/server/Server.cpp

@@ -276,6 +276,10 @@ int Server::main(const std::vector<std::string> & /*args*/)
     if (config().has("max_table_size_to_drop"))
         global_context->setMaxTableSizeToDrop(config().getUInt64("max_table_size_to_drop"));
 
+    /// Setup protection to avoid accidental DROP for big tables (that are greater than 50 GB by default)
+    if (config().has("max_partition_size_to_drop"))
+        global_context->setMaxPartitionSizeToDropAttachReplace(config().getUInt64("max_partition_size_to_drop"));
+
     /// Size of cache for uncompressed blocks. Zero means disabled.
     size_t uncompressed_cache_size = config().getUInt64("uncompressed_cache_size", 0);
     if (uncompressed_cache_size)

dbms/src/Common/ErrorCodes.cpp

@@ -377,6 +377,7 @@ namespace ErrorCodes
     extern const int CANNOT_STAT = 400;
     extern const int FEATURE_IS_NOT_ENABLED_AT_BUILD_TIME = 401;
     extern const int CANNOT_IOSETUP = 402;
+    extern const int PARTITION_SIZE_EXCEEDS_MAX_DROP_SIZE_LIMIT = 403;
 
 
     extern const int KEEPER_EXCEPTION = 999;

dbms/src/Interpreters/Context.cpp

@@ -80,7 +80,7 @@ namespace ErrorCodes
     extern const int THERE_IS_NO_QUERY;
     extern const int NO_ELEMENTS_IN_CONFIG;
     extern const int DDL_GUARD_IS_ACTIVE;
-    extern const int TABLE_SIZE_EXCEEDS_MAX_DROP_SIZE_LIMIT;
+    extern const int PARTITION_SIZE_EXCEEDS_MAX_DROP_SIZE_LIMIT;
     extern const int SESSION_NOT_FOUND;
     extern const int SESSION_IS_LOCKED;
     extern const int CANNOT_GET_CREATE_TABLE_QUERY;
@@ -140,6 +140,7 @@ struct ContextShared
     mutable std::unique_ptr<CompressionSettingsSelector> compression_settings_selector;
     std::unique_ptr<MergeTreeSettings> merge_tree_settings; /// Settings of MergeTree* engines.
     size_t max_table_size_to_drop = 50000000000lu;          /// Protects MergeTree tables from accidental DROP (50GB by default)
+    size_t max_partition_size_to_drop = 50000000000lu;      /// Protects MergeTree partitions from accidental DROP (50GB by default)
     String format_schema_path;                              /// Path to a directory that contains schema files used by input formats.
     ActionLocksManagerPtr action_locks_manager;             /// Set of storages' action lockers
 
@@ -1652,6 +1653,34 @@ void Context::checkTableCanBeDropped(const String & database, const String & tab
 }
 
 
+void Context::setMaxPartitionSizeToDropAttachReplace(size_t max_size)
+{
+    // Is initialized at server startup
+    shared->max_partition_size_to_drop = max_size;
+}
+
+
+void Context::checkPartitionCanBeDroppedAttachReplace(const String & database, const String & table, size_t partition_size)
+{
+    size_t max_partition_size_to_drop = shared->max_partition_size_to_drop;
+    if (!max_partition_size_to_drop || partition_size <= max_partition_size_to_drop)
+        return;
+    
+    String partition_size_str = formatReadableSizeWithDecimalSuffix(partition_size);
+    String max_partition_size_to_drop_str = formatReadableSizeWithDecimalSuffix(max_partition_size_to_drop);
+    std::stringstream ostr;
+
+    ostr << "Partition in table " << backQuoteIfNeed(database) << "." << backQuoteIfNeed(table) << " was not dropped.\n"
+         << "Reason:\n"
+         << "1. Partition size (" << partition_size_str << ") is greater than max_table_size_to_drop (" << max_partition_size_to_drop_str << ")\n";
+
+    ostr << "How to fix this:\n"
+         << "1. Either increase (or set to zero) max_table_size_to_drop in server config and restart ClickHouse\n";
+
+    throw Exception(ostr.str(), ErrorCodes::PARTITION_SIZE_EXCEEDS_MAX_DROP_SIZE_LIMIT);
+}
+
+
 BlockInputStreamPtr Context::getInputFormat(const String & name, ReadBuffer & buf, const Block & sample, size_t max_block_size) const
 {
     return FormatFactory::instance().getInput(name, buf, sample, *this, max_block_size);

dbms/src/Interpreters/Context.h

@@ -368,6 +368,10 @@ public:
     void setMaxTableSizeToDrop(size_t max_size);
     void checkTableCanBeDropped(const String & database, const String & table, size_t table_size);
 
+    /// Prevents DROP PARTITION if its size is greater than max_size (50GB by default, max_size=0 turn off this check)
+    void setMaxPartitionSizeToDropAttachReplace(size_t max_size);
+    void checkPartitionCanBeDroppedAttachReplace(const String & database, const String & table, size_t partition_size);
+
     /// Lets you select the compression settings according to the conditions described in the configuration file.
     CompressionSettings chooseCompressionSettings(size_t part_size, double part_size_ratio) const;
 

dbms/src/Interpreters/InterpreterAlterQuery.cpp

@@ -59,18 +59,23 @@ BlockIO InterpreterAlterQuery::execute()
         switch (command.type)
         {
             case PartitionCommand::DROP_PARTITION:
-                table->dropPartition(query_ptr, command.partition, command.detach, context);
+                if (table->checkPartitionCanBeDroppedAttachReplace(command.partition))
+                    table->dropPartition(query_ptr, command.partition, command.detach, context);
                 break;
 
             case PartitionCommand::ATTACH_PARTITION:
-                table->attachPartition(command.partition, command.part, context);
+                if (table->checkPartitionCanBeDroppedAttachReplace(command.partition))
+                    table->attachPartition(command.partition, command.part, context);
                 break;
 
             case PartitionCommand::REPLACE_PARTITION:
                 {
-                    String from_database = command.from_database.empty() ? context.getCurrentDatabase() : command.from_database;
-                    auto from_storage = context.getTable(from_database, command.from_table);
-                    table->replacePartitionFrom(from_storage, command.partition, command.replace, context);
+                    if (table->checkPartitionCanBeDroppedAttachReplace(command.partition))
+                    {
+                        String from_database = command.from_database.empty() ? context.getCurrentDatabase() : command.from_database;
+                        auto from_storage = context.getTable(from_database, command.from_table);
+                        table->replacePartitionFrom(from_storage, command.partition, command.replace, context);
+                    }
                 }
                 break;
 

dbms/src/Storages/IStorage.h

@@ -329,6 +329,11 @@ public:
     /// Otherwise - throws an exception with detailed information or returns false
     virtual bool checkTableCanBeDropped() const { return true; }
 
+    /// Checks that Partition could be dropped right now
+    /// If it can - returns true
+    /// Otherwise - throws an exception with detailed information or returns false
+    virtual bool checkPartitionCanBeDroppedAttachReplace(const ASTPtr & /*partition*/) { return true; }
+
     /** Notify engine about updated dependencies for this storage. */
     virtual void updateDependencies() {}
 

dbms/src/Storages/StorageMaterializedView.cpp

@@ -293,6 +293,18 @@ bool StorageMaterializedView::checkTableCanBeDropped() const
     return target_table->checkTableCanBeDropped();
 }
 
+bool StorageMaterializedView::checkPartitionCanBeDroppedAttachReplace(const ASTPtr & partition)
+{
+    /// Don't drop the partition in target table if it was created manually via 'TO inner_table' statement
+    if (!has_inner_table)
+        return true;
+
+    auto target_table = tryGetTargetTable();
+    if (!target_table)
+        return true;
+
+    return target_table->checkPartitionCanBeDroppedAttachReplace(partition);
+}
 
 void registerStorageMaterializedView(StorageFactory & factory)
 {

dbms/src/Storages/StorageMaterializedView.h

@@ -41,7 +41,9 @@ public:
     void freezePartition(const ASTPtr & partition, const String & with_name, const Context & context) override;
 
     void shutdown() override;
+
     bool checkTableCanBeDropped() const override;
+    bool checkPartitionCanBeDroppedAttachReplace(const ASTPtr & partition) override;
 
     BlockInputStreams read(
         const Names & column_names,

dbms/src/Storages/StorageMergeTree.cpp

@@ -129,6 +129,20 @@ bool StorageMergeTree::checkTableCanBeDropped() const
     return true;
 }
 
+bool StorageMergeTree::checkPartitionCanBeDroppedAttachReplace(const ASTPtr & partition)
+{
+    const_cast<MergeTreeData &>(getData()).recalculateColumnSizes();
+    
+    const String partition_id = data.getPartitionIDFromQuery(partition, context);
+    auto parts_to_remove = data.getDataPartsVectorInPartition(MergeTreeDataPartState::Committed, partition_id);
+    
+    for (const auto & part : parts_to_remove)
+    {
+        context.checkPartitionCanBeDroppedAttachReplace(database_name, table_name, part->getTotalColumnsSize().data_compressed);
+    }
+    return true;
+}
+
 void StorageMergeTree::drop()
 {
     shutdown();

dbms/src/Storages/StorageMergeTree.h

@@ -85,6 +85,8 @@ public:
 
     bool checkTableCanBeDropped() const override;
 
+    bool checkPartitionCanBeDroppedAttachReplace(const ASTPtr & partition) override;
+
     ActionLock getActionLock(StorageActionBlockType action_type) override;
 
     MergeTreeData & getData() { return data; }

dbms/src/Storages/StorageReplicatedMergeTree.cpp

@@ -3352,6 +3352,21 @@ bool StorageReplicatedMergeTree::checkTableCanBeDropped() const
 }
 
 
+bool StorageReplicatedMergeTree::checkPartitionCanBeDroppedAttachReplace(const ASTPtr & partition)
+{
+    const_cast<MergeTreeData &>(getData()).recalculateColumnSizes();
+    
+    const String partition_id = data.getPartitionIDFromQuery(partition, context);
+    auto parts_to_remove = data.getDataPartsVectorInPartition(MergeTreeDataPartState::Committed, partition_id);
+    
+    for (const auto & part : parts_to_remove)
+    {
+        context.checkPartitionCanBeDroppedAttachReplace(database_name, table_name, part->getTotalColumnsSize().data_compressed);
+    }
+    return true;
+}
+
+
 void StorageReplicatedMergeTree::drop()
 {
     {

dbms/src/Storages/StorageReplicatedMergeTree.h

@@ -140,6 +140,8 @@ public:
 
     bool checkTableCanBeDropped() const override;
 
+    bool checkPartitionCanBeDroppedAttachReplace(const ASTPtr & partition) override;
+
     ActionLock getActionLock(StorageActionBlockType action_type) override;
 
     /// Wait when replication queue size becomes less or equal than queue_size