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

@@ -141,6 +141,7 @@ class IColumn;
     \
     M(Bool, force_index_by_date, 0, "Throw an exception if there is a partition key in a table, and it is not used.", 0) \
     M(Bool, force_primary_key, 0, "Throw an exception if there is primary key in a table, and it is not used.", 0) \
+    M(Bool, force_max_partition_limit, 0, "Throw an exception if max_partitions_to_read is violated.", 0) \
     M(String, force_data_skipping_indices, "", "Comma separated list of strings or literals with the name of the data skipping indices that should be used during query execution, otherwise an exception will be thrown.", 0) \
     \
     M(Float, max_streams_to_max_threads_ratio, 1, "Allows you to use more sources than the number of threads - to more evenly distribute work across threads. It is assumed that this is a temporary solution, since it will be possible in the future to make the number of sources equal to the number of threads, but for each source to dynamically select available work for itself.", 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,20 @@ QueryPlanPtr MergeTreeDataSelectExecutor::readFromParts(
     if (parts_with_ranges.empty())
         return std::make_unique<QueryPlan>();
 
+    auto max_partitions_to_read = data.getSettings()->max_partitions_to_read;
+    if (settings.force_max_partition_limit && 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.", 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-03	1	2
+2021-01-04	3	4
+2021-01-03	1	2
+2021-01-04	3	4
+2021-01-03	1	2
+2021-01-04	3	4

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 (yesterday(), 1, 2), (today(), 3, 4);
+
+select * from p order by i; -- default no limit
+
+select * from p order by i settings force_max_partition_limit = 0;
+
+select * from p order by i settings force_max_partition_limit = 1; -- { serverError 565 }
+
+alter table p modify setting max_partitions_to_read = 2;
+
+select * from p order by i settings force_max_partition_limit = 1;
+
+drop table if exists p;