Merge pull request #18712 from amosbird/limitpartition

add max_partitions_to_read setting

src/Common/ErrorCodes.cpp

@@ -531,6 +531,7 @@
     M(562, TLD_LIST_NOT_FOUND) \
     M(563, CANNOT_READ_MAP_FROM_TEXT) \
     M(564, INTERSERVER_SCHEME_DOESNT_MATCH) \
+    M(565, TOO_MANY_PARTITIONS) \
     \
     M(999, KEEPER_EXCEPTION) \
     M(1000, POCO_EXCEPTION) \

src/Core/Settings.h

@@ -353,6 +353,7 @@ class IColumn;
     M(Bool, allow_introspection_functions, false, "Allow functions for introspection of ELF and DWARF for query profiling. These functions are slow and may impose security considerations.", 0) \
     \
     M(UInt64, max_partitions_per_insert_block, 100, "Limit maximum number of partitions in single INSERTed block. Zero means unlimited. Throw exception if the block contains too many partitions. This setting is a safety threshold, because using large number of partitions is a common misconception.", 0) \
+    M(UInt64, max_partitions_to_read, 0, "Limit the max number of partitions that can be accessed in one query. 0 means unlimited.", 0) \
     M(Bool, check_query_single_value_result, true, "Return check query result as single 1/0 value", 0) \
     M(Bool, allow_drop_detached, false, "Allow ALTER TABLE ... DROP DETACHED PART[ITION] ... queries", 0) \
     \

src/Storages/MergeTree/MergeTreeDataSelectExecutor.cpp

@@ -59,6 +59,7 @@ namespace ErrorCodes
     extern const int ARGUMENT_OUT_OF_BOUND;
     extern const int TOO_MANY_ROWS;
     extern const int CANNOT_PARSE_TEXT;
+    extern const int TOO_MANY_PARTITIONS;
 }
 
 
@@ -706,6 +707,21 @@ QueryPlanPtr MergeTreeDataSelectExecutor::readFromParts(
     if (parts_with_ranges.empty())
         return std::make_unique<QueryPlan>();
 
+    auto max_partitions_to_read
+        = settings.max_partitions_to_read.changed ? settings.max_partitions_to_read : data.getSettings()->max_partitions_to_read;
+    if (max_partitions_to_read)
+    {
+        std::set<String> partitions;
+        for (auto & part_with_ranges : parts_with_ranges)
+            partitions.insert(part_with_ranges.data_part->info.partition_id);
+        if (partitions.size() > max_partitions_to_read)
+            throw Exception(
+                ErrorCodes::TOO_MANY_PARTITIONS,
+                "Too many partitions to read. Current {}, max {}",
+                partitions.size(),
+                max_partitions_to_read);
+    }
+
     ProfileEvents::increment(ProfileEvents::SelectedParts, parts_with_ranges.size());
     ProfileEvents::increment(ProfileEvents::SelectedRanges, sum_ranges);
     ProfileEvents::increment(ProfileEvents::SelectedMarks, sum_marks);

src/Storages/MergeTree/MergeTreeSettings.h

@@ -110,6 +110,7 @@ struct Settings;
     M(Bool, allow_nullable_key, false, "Allow Nullable types as primary keys.", 0) \
     M(Bool, remove_empty_parts, true, "Remove empty parts after they were pruned by TTL, mutation, or collapsing merge algorithm", 0) \
     M(Bool, assign_part_uuids, false, "Generate UUIDs for parts. Before enabling check that all replicas support new format.", 0) \
+    M(UInt64, max_partitions_to_read, 0, "Limit the max number of partitions that can be accessed in one query. 0 means unlimited. This setting is the default that can be overridden by the query-level setting with the same name.", 0) \
     \
     /** Obsolete settings. Kept for backward compatibility only. */ \
     M(UInt64, min_relative_delay_to_yield_leadership, 120, "Obsolete setting, does nothing.", 0) \

tests/queries/0_stateless/01632_max_partitions_to_read.reference

@@ -0,0 +1,6 @@
+2021-01-01	1	2
+2021-01-02	4	5
+2021-01-01	1	2
+2021-01-02	4	5
+2021-01-01	1	2
+2021-01-02	4	5

tests/queries/0_stateless/01632_max_partitions_to_read.sql

@@ -0,0 +1,17 @@
+drop table if exists p;
+
+create table p(d Date, i int, j int) engine MergeTree partition by d order by i settings max_partitions_to_read = 1;
+
+insert into p values ('2021-01-01', 1, 2), ('2021-01-02', 4, 5);
+
+select * from p order by i; -- { serverError 565 }
+
+select * from p order by i settings max_partitions_to_read = 2;
+
+select * from p order by i settings max_partitions_to_read = 0; -- unlimited
+
+alter table p modify setting max_partitions_to_read = 2;
+
+select * from p order by i;
+
+drop table if exists p;