Add guardrail for number of partition keys on IN queries

patch by Krishna Vadali; reviewed by Andrés de la Peña and Berenguer Blasi for CASSANDRA-17186
This commit is contained in:
tejavadali 2022-02-23 12:42:42 +00:00 committed by Andrés de la Peña
parent ab7ee50013
commit 7ea0a4aa41
9 changed files with 189 additions and 4 deletions

View File

@ -1,4 +1,5 @@
4.1
* Add guardrail for number of partition keys on IN queries (CASSANDRA-17186)
* update Python test framework from nose to pytest (CASSANDRA-17293)
* Fix improper CDC commit log segments deletion in non-blocking mode (CASSANDRA-17233)
* Add support for string concatenations through the + operator (CASSANDRA-17190)

View File

@ -1616,3 +1616,7 @@ drop_compact_storage_enabled: false
# Guardrail to allow/disallow list operations that require read before write, i.e. setting list element by index and
# removing list elements by either index or value. Defaults to true.
# read_before_write_list_operations_enabled: true
# Guardrail to warn or fail when querying with an IN restriction selecting more partition keys than threshold.
# The two thresholds default to -1 to disable.
# partition_keys_in_select_warn_threshold: -1
# partition_keys_in_select_fail_threshold: -1

View File

@ -738,6 +738,8 @@ public class Config
public volatile int materialized_views_per_table_fail_threshold = DISABLED_GUARDRAIL;
public volatile int page_size_warn_threshold = DISABLED_GUARDRAIL;
public volatile int page_size_fail_threshold = DISABLED_GUARDRAIL;
public volatile int partition_keys_in_select_warn_threshold = DISABLED_GUARDRAIL;
public volatile int partition_keys_in_select_fail_threshold = DISABLED_GUARDRAIL;
public volatile Set<String> table_properties_ignored = Collections.emptySet();
public volatile Set<String> table_properties_disallowed = Collections.emptySet();
public volatile boolean user_timestamps_enabled = true;

View File

@ -67,6 +67,8 @@ public class GuardrailsOptions implements GuardrailsConfig
config.table_properties_ignored = validateTableProperties(config.table_properties_ignored, "table_properties_ignored");
config.table_properties_disallowed = validateTableProperties(config.table_properties_disallowed, "table_properties_disallowed");
validateIntThreshold(config.page_size_warn_threshold, config.page_size_fail_threshold, "page_size");
validateIntThreshold(config.partition_keys_in_select_warn_threshold,
config.partition_keys_in_select_fail_threshold, "partition_keys_in_select");
}
@Override
@ -194,6 +196,31 @@ public class GuardrailsOptions implements GuardrailsConfig
return config.materialized_views_per_table_warn_threshold;
}
@Override
public int getPartitionKeysInSelectWarnThreshold()
{
return config.partition_keys_in_select_warn_threshold;
}
@Override
public int getPartitionKeysInSelectFailThreshold()
{
return config.partition_keys_in_select_fail_threshold;
}
public void setPartitionKeysInSelectThreshold(int warn, int fail)
{
validateIntThreshold(warn, fail, "partition_keys_in_select");
updatePropertyWithLogging("partition_keys_in_select_warn_threshold",
warn,
() -> config.partition_keys_in_select_warn_threshold,
x -> config.partition_keys_in_select_warn_threshold = x);
updatePropertyWithLogging("partition_keys_in_select_fail_threshold",
fail,
() -> config.partition_keys_in_select_fail_threshold,
x -> config.partition_keys_in_select_fail_threshold = x);
}
@Override
public int getMaterializedViewsPerTableFailThreshold()
{

View File

@ -247,7 +247,8 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
int pageSize = options.getPageSize();
Selectors selectors = selection.newSelectors(options);
ReadQuery query = getQuery(options, selectors.getColumnFilter(), nowInSec, userLimit, userPerPartitionLimit, pageSize);
ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(),
nowInSec, userLimit, userPerPartitionLimit, pageSize);
if (options.isTrackWarningsEnabled())
query.trackWarnings();
@ -271,6 +272,7 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
{
Selectors selectors = selection.newSelectors(options);
return getQuery(options,
ClientState.forInternalCalls(),
selectors.getColumnFilter(),
nowInSec,
getLimit(options),
@ -279,6 +281,7 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
}
public ReadQuery getQuery(QueryOptions options,
ClientState state,
ColumnFilter columnFilter,
int nowInSec,
int userLimit,
@ -292,7 +295,7 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
if (isPartitionRangeQuery)
return getRangeCommand(options, columnFilter, limit, nowInSec);
return getSliceCommands(options, columnFilter, limit, nowInSec);
return getSliceCommands(options, state, columnFilter, limit, nowInSec);
}
private ResultMessage.Rows execute(ReadQuery query,
@ -452,7 +455,8 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
int pageSize = options.getPageSize();
Selectors selectors = selection.newSelectors(options);
ReadQuery query = getQuery(options, selectors.getColumnFilter(), nowInSec, userLimit, userPerPartitionLimit, pageSize);
ReadQuery query = getQuery(options, state.getClientState(), selectors.getColumnFilter(), nowInSec, userLimit,
userPerPartitionLimit, pageSize);
try (ReadExecutionController executionController = query.executionController())
{
@ -521,12 +525,18 @@ public class SelectStatement implements CQLStatement.SingleKeyspaceCqlStatement
return restrictions;
}
private ReadQuery getSliceCommands(QueryOptions options, ColumnFilter columnFilter, DataLimits limit, int nowInSec)
private ReadQuery getSliceCommands(QueryOptions options, ClientState state, ColumnFilter columnFilter,
DataLimits limit, int nowInSec)
{
Collection<ByteBuffer> keys = restrictions.getPartitionKeys(options);
if (keys.isEmpty())
return ReadQuery.empty(table);
if (restrictions.keyIsInRelation())
{
Guardrails.partitionKeysInSelect.guard(keys.size(), table.name, state);
}
ClusteringIndexFilter filter = makeClusteringIndexFilter(options, columnFilter);
if (filter == null || filter.isEmpty(table.comparator))
return ReadQuery.empty(table);

View File

@ -127,6 +127,20 @@ public final class Guardrails implements GuardrailsMBean
: format("Aborting query for table %s, page size %s exceeds fail threshold of %s.",
what, value, threshold));
/**
* Guardrail on the number of partition keys in the IN clause.
*/
public static final Threshold partitionKeysInSelect =
new Threshold(state -> CONFIG_PROVIDER.getOrCreate(state).getPartitionKeysInSelectWarnThreshold(),
state -> CONFIG_PROVIDER.getOrCreate(state).getPartitionKeysInSelectFailThreshold(),
(isWarning, what, value, threshold) ->
isWarning ? format("Query with partition keys in IN clause on table %s, with number of " +
"partition keys %s exceeds warning threshold of %s.",
what, value, threshold)
: format("Aborting query with partition keys in IN clause on table %s, " +
"number of partition keys %s exceeds fail threshold of %s.",
what, value, threshold));
/**
* Guardrail disabling operations on lists that require read before write.
*/
@ -351,6 +365,24 @@ public final class Guardrails implements GuardrailsMBean
DEFAULT_CONFIG.setReadBeforeWriteListOperationsEnabled(enabled);
}
@Override
public void setPartitionKeysInSelectThreshold(int warn, int fail)
{
DEFAULT_CONFIG.setPartitionKeysInSelectThreshold(warn, fail);
}
@Override
public int getPartitionKeysInSelectWarnThreshold()
{
return DEFAULT_CONFIG.getPartitionKeysInSelectWarnThreshold();
}
@Override
public int getPartitionKeysInSelectFailThreshold()
{
return DEFAULT_CONFIG.getPartitionKeysInSelectFailThreshold();
}
private static String toCSV(Set<String> values)
{
return values == null ? "" : String.join(",", values);

View File

@ -95,6 +95,17 @@ public interface GuardrailsConfig
*/
int getMaterializedViewsPerTableWarnThreshold();
/**
* @return The threshold to warn when partition keys in select more than threshold.
*/
int getPartitionKeysInSelectWarnThreshold();
/**
* @return The threshold to fail when partition keys in select more than threshold.
*/
int getPartitionKeysInSelectFailThreshold();
/**
* @return The threshold to fail when creating more materialized views per table than threshold.
*/

View File

@ -220,4 +220,25 @@ public interface GuardrailsMBean
* @param enabled {@code true} if list operations that require read before write are allowed, {@code false} otherwise.
*/
void setReadBeforeWriteListOperationsEnabled(boolean enabled);
/**
* @param warn The threshold to warn when the number of partition keys in a select statement is greater than
* threshold -1 means disabled.
* @param fail The threshold to prevent when the number of partition keys in a select statement is more than
* threshold -1 means disabled.
*/
void setPartitionKeysInSelectThreshold(int warn, int fail);
/**
* @return The threshold to warn when the number of partition keys in a select statement greater than threshold.
* -1 means disabled.
*/
int getPartitionKeysInSelectWarnThreshold();
/**
* @return The threshold to fail when the number of partition keys in a select statement greater than threshold.
* -1 means disabled.
*/
int getPartitionKeysInSelectFailThreshold();
}

View File

@ -0,0 +1,77 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.cassandra.db.guardrails;
import org.junit.Before;
import org.junit.Test;
public class GuardrailPartitionKeysInSelectTest extends ThresholdTester
{
private static final int PARTITION_KEYS_SELECT_WARN_THRESHOLD = 3;
private static final int PARTITION_KEYS_SELECT_FAIL_THRESHOLD = 5;
private String tableName;
public GuardrailPartitionKeysInSelectTest()
{
super(PARTITION_KEYS_SELECT_WARN_THRESHOLD,
PARTITION_KEYS_SELECT_FAIL_THRESHOLD,
"partition_keys_in_select",
Guardrails::setPartitionKeysInSelectThreshold,
Guardrails::getPartitionKeysInSelectWarnThreshold,
Guardrails::getPartitionKeysInSelectFailThreshold);
}
@Before
public void setupTest()
{
tableName = createTable("CREATE TABLE %s (k INT, c INT, v TEXT, PRIMARY KEY(k, c))");
}
@Test
public void testSelectStatementAgainstInClausePartitionKeys() throws Throwable
{
assertValid("SELECT k, c, v FROM %s WHERE k=10");
assertValid("SELECT k, c, v FROM %s WHERE k IN (2, 3)");
assertValid("SELECT k, c, v FROM %s WHERE k = 2 and c IN (2, 3, 4, 5, 6, 7)");
assertWarns(String.format("Query with partition keys in IN clause on table %s, with " +
"number of partition keys 4 exceeds warning threshold of 3.", tableName),
"SELECT k, c, v FROM %s WHERE k IN (2, 3, 4, 5)");
assertFails(String.format("Aborting query with partition keys in IN clause on table %s, " +
"number of partition keys 6 exceeds fail threshold of 5.", tableName) ,
"SELECT k, c, v FROM %s WHERE k IN (2, 3, 4, 5, 6, 7)"
);
}
@Test
public void testExcludedUsers() throws Throwable
{
testExcludedUsers(() -> "SELECT k, c, v FROM %s WHERE k IN (2, 3, 4, 5)",
() -> "SELECT k, c, v FROM %s WHERE k IN (2, 3, 4, 5, 6, 7)");
}
@Override
protected long currentValue()
{
throw new UnsupportedOperationException();
}
}