diff --git a/src/Common/ErrorCodes.cpp b/src/Common/ErrorCodes.cpp index 8b14e17a138..1ceb9df011d 100644 --- a/src/Common/ErrorCodes.cpp +++ b/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) \ diff --git a/src/Core/Settings.h b/src/Core/Settings.h index 00d4682332d..e2aa4e2f2d8 100644 --- a/src/Core/Settings.h +++ b/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) \ \ diff --git a/src/Storages/MergeTree/MergeTreeDataSelectExecutor.cpp b/src/Storages/MergeTree/MergeTreeDataSelectExecutor.cpp index 98a08abab65..75760e145d3 100644 --- a/src/Storages/MergeTree/MergeTreeDataSelectExecutor.cpp +++ b/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(); + 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 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); diff --git a/src/Storages/MergeTree/MergeTreeSettings.h b/src/Storages/MergeTree/MergeTreeSettings.h index 9b344d19f8b..97c6114a0be 100644 --- a/src/Storages/MergeTree/MergeTreeSettings.h +++ b/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) \ diff --git a/tests/queries/0_stateless/01632_max_partitions_to_read.reference b/tests/queries/0_stateless/01632_max_partitions_to_read.reference new file mode 100644 index 00000000000..ea2526e1301 --- /dev/null +++ b/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 diff --git a/tests/queries/0_stateless/01632_max_partitions_to_read.sql b/tests/queries/0_stateless/01632_max_partitions_to_read.sql new file mode 100644 index 00000000000..b91405569bc --- /dev/null +++ b/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;