diff --git a/docs/api-reference/dynamic-configuration-api.md b/docs/api-reference/dynamic-configuration-api.md index 259e50c88008..57774a9c279a 100644 --- a/docs/api-reference/dynamic-configuration-api.md +++ b/docs/api-reference/dynamic-configuration-api.md @@ -371,6 +371,7 @@ Host: http://ROUTER_IP:ROUTER_PORT { "queryBlocklist": [ { + "type": "default", "ruleName": "block-expensive-scans", "dataSources": ["large_table"], "queryTypes": ["scan"] @@ -440,11 +441,13 @@ curl -X POST "http://ROUTER_IP:ROUTER_PORT/druid/coordinator/v1/broker/config" \ -d '{ "queryBlocklist": [ { + "type": "default", "ruleName": "block-expensive-scans", "dataSources": ["large_table", "huge_dataset"], "queryTypes": ["scan"] }, { + "type": "default", "ruleName": "block-debug-queries", "contextMatches": { "debug": "true" @@ -481,11 +484,13 @@ X-Druid-Comment: Add query blocklist rules and set default context { "queryBlocklist": [ { + "type": "default", "ruleName": "block-expensive-scans", "dataSources": ["large_table", "huge_dataset"], "queryTypes": ["scan"] }, { + "type": "default", "ruleName": "block-debug-queries", "contextMatches": { "debug": "true" @@ -529,15 +534,34 @@ The following table shows the dynamic configuration properties for the Broker. Query blocklist rules allow you to block specific queries based on datasource, query type, and/or query context parameters. This feature is useful for preventing expensive or problematic queries from impacting cluster performance. -Each rule in the `queryBlocklist` array is a JSON object with the following properties: +Each rule in the `queryBlocklist` array is a JSON object. The `type` field selects the rule implementation. If omitted, it defaults to `"default"`. + +##### `default` type + +The built-in rule type. Blocks queries by matching on datasource, query type, and/or query context parameters. |Property|Description|Required|Default| |--------|-----------|--------|-------| +|`type`|Rule type identifier. Not required — rules without a `"type"` field are treated as `"default"` for backwards compatibility with existing configurations.|No|`"default"`| |`ruleName`|Unique name identifying this blocklist rule. Used in error messages when queries are blocked.|Yes|N/A| |`dataSources`|List of datasource names to match. A query matches if it references any datasource in this list.|No|Matches all datasources| |`queryTypes`|List of query types to match (e.g., `scan`, `timeseries`, `groupBy`, `topN`). A query matches if its type is in this list.|No|Matches all query types| |`contextMatches`|Map of query context parameter key-value pairs to match. A query matches if all specified context parameters match the provided values (case-sensitive string comparison).|No|Matches all contexts| +##### Custom types (extensions) + +Extensions can register additional rule types by adding Jackson subtypes to the `QueryBlocklistRule` interface. A custom rule is selected by setting `"type"` to its registered name: + +```json +{ + "type": "myCustomRule", + "ruleName": "rate-limit-heavy-users", + ... +} +``` + +If a rule references a type whose extension is not loaded, deserialization fails with an error rather than silently falling back to the default type. + **Rule matching behavior:** - A query must match ALL specified criteria within a rule (AND logic) to be blocked by that rule @@ -546,6 +570,8 @@ Each rule in the `queryBlocklist` array is a JSON object with the following prop - At least one criterion must be specified per rule to prevent accidentally blocking all queries - A query is blocked if it matches ANY rule in the blocklist (OR logic between rules) +> **Note:** The `"type"` field is not required for `"default"` rules (existing configs without it continue to work), but including it explicitly is recommended for clarity. + **Error response:** When a query is blocked, the Broker returns an HTTP 403 error with a message indicating the query ID and the rule that blocked it: @@ -676,7 +702,7 @@ Host: http://ROUTER_IP:ROUTER_PORT "comment": "Add query blocklist rules", "ip": "127.0.0.1" }, - "payload": "{\"queryBlocklist\":[{\"ruleName\":\"block-expensive-scans\",\"dataSources\":[\"large_table\"],\"queryTypes\":[\"scan\"]}],\"queryContext\":{\"priority\":0,\"timeout\":300000},\"perSegmentTimeoutConfig\":{\"large_table\":{\"perSegmentTimeoutMs\":10000,\"monitorOnly\":false}}}", + "payload": "{\"queryBlocklist\":[{\"type\":\"default\",\"ruleName\":\"block-expensive-scans\",\"dataSources\":[\"large_table\"],\"queryTypes\":[\"scan\"]}],\"queryContext\":{\"priority\":0,\"timeout\":300000},\"perSegmentTimeoutConfig\":{\"large_table\":{\"perSegmentTimeoutMs\":10000,\"monitorOnly\":false}}}", "auditTime": "2024-03-06T12:00:00.000Z" } ] diff --git a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/EmbeddedBrokerDynamicConfigTest.java b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/EmbeddedBrokerDynamicConfigTest.java index 34020229434e..ac48e2c53466 100644 --- a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/EmbeddedBrokerDynamicConfigTest.java +++ b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/server/EmbeddedBrokerDynamicConfigTest.java @@ -24,6 +24,7 @@ import org.apache.druid.common.utils.IdUtils; import org.apache.druid.indexing.common.task.TaskBuilder; import org.apache.druid.query.QueryContext; +import org.apache.druid.server.DefaultQueryBlocklistRule; import org.apache.druid.server.QueryBlocklistRule; import org.apache.druid.server.broker.BrokerDynamicConfig; import org.apache.druid.server.http.BrokerDynamicConfigSyncer; @@ -102,7 +103,7 @@ public void testQueryBlocklistBlocksMatchingQueries() Assertions.assertFalse(initialResult.isBlank()); // Apply blocklist rule that matches all queries on this datasource - QueryBlocklistRule blockRule = new QueryBlocklistRule( + QueryBlocklistRule blockRule = new DefaultQueryBlocklistRule( "block-test-datasource", Set.of(dataSource), null, diff --git a/server/src/main/java/org/apache/druid/server/DefaultQueryBlocklistRule.java b/server/src/main/java/org/apache/druid/server/DefaultQueryBlocklistRule.java new file mode 100644 index 000000000000..ae1d8662d28d --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/DefaultQueryBlocklistRule.java @@ -0,0 +1,172 @@ +/* + * 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.druid.server; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; +import com.google.common.base.Strings; +import com.google.common.collect.Sets; +import org.apache.druid.query.Query; + +import javax.annotation.Nullable; +import java.util.Map; +import java.util.Objects; +import java.util.Set; + +/** + * Default {@link QueryBlocklistRule} implementation. A query matches if ALL specified criteria + * match (AND logic). Null or empty criteria are wildcards (match everything). + * + *

At least one criterion must be non-empty to prevent accidentally blocking all queries. + * Deserializes from JSON with no {@code "type"} field for backwards compatibility. + */ +public class DefaultQueryBlocklistRule implements QueryBlocklistRule +{ + private final String ruleName; + @Nullable + private final Set dataSources; + @Nullable + private final Set queryTypes; + @Nullable + private final Map contextMatches; + + private final boolean hasDataSourceCriteria; + private final boolean hasQueryTypeCriteria; + private final boolean hasContextCriteria; + + @JsonCreator + public DefaultQueryBlocklistRule( + @JsonProperty("ruleName") String ruleName, + @JsonProperty("dataSources") @Nullable Set dataSources, + @JsonProperty("queryTypes") @Nullable Set queryTypes, + @JsonProperty("contextMatches") @Nullable Map contextMatches + ) + { + Preconditions.checkArgument( + !Strings.isNullOrEmpty(ruleName), + "ruleName must not be null or empty" + ); + + this.hasDataSourceCriteria = dataSources != null && !dataSources.isEmpty(); + this.hasQueryTypeCriteria = queryTypes != null && !queryTypes.isEmpty(); + this.hasContextCriteria = contextMatches != null && !contextMatches.isEmpty(); + + Preconditions.checkArgument( + hasDataSourceCriteria || hasQueryTypeCriteria || hasContextCriteria, + "At least one criterion (dataSources, queryTypes, or contextMatches) must be specified. " + + "A rule with all null/empty criteria would block ALL queries." + ); + + this.ruleName = ruleName; + this.dataSources = dataSources; + this.queryTypes = queryTypes; + this.contextMatches = contextMatches; + } + + @Override + @JsonProperty + public String getRuleName() + { + return ruleName; + } + + @JsonProperty + @Nullable + public Set getDataSources() + { + return dataSources; + } + + @JsonProperty + @Nullable + public Set getQueryTypes() + { + return queryTypes; + } + + @JsonProperty + @Nullable + public Map getContextMatches() + { + return contextMatches; + } + + @Override + public boolean matches(Query query) + { + if (hasDataSourceCriteria) { + Set queryDatasources = query.getDataSource().getTableNames(); + if (Sets.intersection(dataSources, queryDatasources).isEmpty()) { + return false; + } + } + + if (hasQueryTypeCriteria) { + if (!queryTypes.contains(query.getType())) { + return false; + } + } + + if (hasContextCriteria) { + for (Map.Entry entry : contextMatches.entrySet()) { + Object contextValue = query.getContext().get(entry.getKey()); + if (contextValue == null || !entry.getValue().equals(String.valueOf(contextValue))) { + return false; + } + } + } + + return true; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + DefaultQueryBlocklistRule that = (DefaultQueryBlocklistRule) o; + return Objects.equals(ruleName, that.ruleName) + && Objects.equals(dataSources, that.dataSources) + && Objects.equals(queryTypes, that.queryTypes) + && Objects.equals(contextMatches, that.contextMatches); + } + + @Override + public int hashCode() + { + return Objects.hash(ruleName, dataSources, queryTypes, contextMatches); + } + + @Override + public String toString() + { + return "DefaultQueryBlocklistRule{" + + "ruleName='" + ruleName + '\'' + + ", dataSources=" + dataSources + + ", queryTypes=" + queryTypes + + ", contextMatches=" + contextMatches + + '}'; + } +} diff --git a/server/src/main/java/org/apache/druid/server/QueryBlocklistRule.java b/server/src/main/java/org/apache/druid/server/QueryBlocklistRule.java index 1242a11bff32..55b20c262a3e 100644 --- a/server/src/main/java/org/apache/druid/server/QueryBlocklistRule.java +++ b/server/src/main/java/org/apache/druid/server/QueryBlocklistRule.java @@ -19,158 +19,37 @@ package org.apache.druid.server; -import com.fasterxml.jackson.annotation.JsonCreator; -import com.fasterxml.jackson.annotation.JsonProperty; -import com.google.common.base.Preconditions; -import com.google.common.base.Strings; -import com.google.common.collect.Sets; +import com.fasterxml.jackson.annotation.JsonSubTypes; +import com.fasterxml.jackson.annotation.JsonTypeInfo; +import com.fasterxml.jackson.databind.annotation.JsonTypeResolver; +import org.apache.druid.jackson.StrictTypeIdResolver; import org.apache.druid.query.Query; -import javax.annotation.Nullable; -import java.util.Map; -import java.util.Objects; -import java.util.Set; - /** - * A rule for matching queries against blocklist criteria. A query matches this rule if ALL - * specified criteria match (AND logic). Null or empty criteria match everything. + * A rule that determines whether a query should be blocked. Implementations define their own + * matching logic and JSON fields. Use {@link DefaultQueryBlocklistRule} for the standard criteria + * (datasources, query types, context). + * + *

Rules with no {@code "type"} field in JSON deserialize as {@link DefaultQueryBlocklistRule} + * for backwards compatibility. An explicit but unrecognized {@code "type"} value will fail + * deserialization rather than silently falling back to the default — this prevents extension + * rules from being misinterpreted when the extension is not loaded. Extensions can register + * additional implementations as Jackson subtypes via {@code SimpleModule.registerSubtypes(...)}. + * + *

Implementations must define {@code equals} and {@code hashCode} so that + * {@link org.apache.druid.server.broker.BrokerDynamicConfig} can detect changes correctly. */ -public class QueryBlocklistRule +@JsonTypeResolver(StrictTypeIdResolver.Builder.class) +@JsonTypeInfo(use = JsonTypeInfo.Id.CUSTOM, property = "type", defaultImpl = DefaultQueryBlocklistRule.class) +@JsonSubTypes({ + @JsonSubTypes.Type(value = DefaultQueryBlocklistRule.class, name = "default") +}) +public interface QueryBlocklistRule { - private final String ruleName; - @Nullable - private final Set dataSources; - @Nullable - private final Set queryTypes; - @Nullable - private final Map contextMatches; - - private final boolean hasDataSourceCriteria; - private final boolean hasQueryTypeCriteria; - private final boolean hasContextCriteria; - - @JsonCreator - public QueryBlocklistRule( - @JsonProperty("ruleName") String ruleName, - @JsonProperty("dataSources") @Nullable Set dataSources, - @JsonProperty("queryTypes") @Nullable Set queryTypes, - @JsonProperty("contextMatches") @Nullable Map contextMatches - ) - { - Preconditions.checkArgument( - !Strings.isNullOrEmpty(ruleName), - "ruleName must not be null or empty" - ); - - // At least one criterion must be specified to prevent accidentally blocking all queries - this.hasDataSourceCriteria = dataSources != null && !dataSources.isEmpty(); - this.hasQueryTypeCriteria = queryTypes != null && !queryTypes.isEmpty(); - this.hasContextCriteria = contextMatches != null && !contextMatches.isEmpty(); - - Preconditions.checkArgument( - hasDataSourceCriteria || hasQueryTypeCriteria || hasContextCriteria, - "At least one criterion (dataSources, queryTypes, or contextMatches) must be specified. " - + "A rule with all null/empty criteria would block ALL queries." - ); - - this.ruleName = ruleName; - this.dataSources = dataSources; - this.queryTypes = queryTypes; - this.contextMatches = contextMatches; - } - - @JsonProperty - public String getRuleName() - { - return ruleName; - } - - @JsonProperty - @Nullable - public Set getDataSources() - { - return dataSources; - } - - @JsonProperty - @Nullable - public Set getQueryTypes() - { - return queryTypes; - } - - @JsonProperty - @Nullable - public Map getContextMatches() - { - return contextMatches; - } + String getRuleName(); /** - * Returns true if the query matches ALL specified criteria (AND logic). - * Null or empty criteria match everything. - * - * @param query the query to check - * @return true if the query matches this rule, false otherwise + * Returns true if the query matches this rule and should be blocked. */ - public boolean matches(Query query) - { - if (hasDataSourceCriteria) { - Set queryDatasources = query.getDataSource().getTableNames(); - if (Sets.intersection(dataSources, queryDatasources).isEmpty()) { - return false; - } - } - - if (hasQueryTypeCriteria) { - if (!queryTypes.contains(query.getType())) { - return false; - } - } - - if (hasContextCriteria) { - for (Map.Entry entry : contextMatches.entrySet()) { - Object contextValue = query.getContext().get(entry.getKey()); - // If the query context doesn't have this key or has a null value, it doesn't match - if (contextValue == null || !entry.getValue().equals(String.valueOf(contextValue))) { - return false; - } - } - } - - return true; - } - - @Override - public boolean equals(Object o) - { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - QueryBlocklistRule that = (QueryBlocklistRule) o; - return Objects.equals(ruleName, that.ruleName) - && Objects.equals(dataSources, that.dataSources) - && Objects.equals(queryTypes, that.queryTypes) - && Objects.equals(contextMatches, that.contextMatches); - } - - @Override - public int hashCode() - { - return Objects.hash(ruleName, dataSources, queryTypes, contextMatches); - } - - @Override - public String toString() - { - return "QueryBlocklistRule{" + - "ruleName='" + ruleName + '\'' + - ", dataSources=" + dataSources + - ", queryTypes=" + queryTypes + - ", contextMatches=" + contextMatches + - '}'; - } + boolean matches(Query query); } diff --git a/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java b/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java index d915ece0c2ed..0b39eb150160 100644 --- a/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java +++ b/server/src/test/java/org/apache/druid/client/coordinator/CoordinatorClientImplTest.java @@ -49,7 +49,7 @@ import org.apache.druid.segment.column.ColumnType; import org.apache.druid.segment.column.RowSignature; import org.apache.druid.segment.metadata.DataSourceInformation; -import org.apache.druid.server.QueryBlocklistRule; +import org.apache.druid.server.DefaultQueryBlocklistRule; import org.apache.druid.server.broker.BrokerDynamicConfig; import org.apache.druid.server.compaction.CompactionStatusResponse; import org.apache.druid.server.coordination.DruidServerMetadata; @@ -840,7 +840,7 @@ public void test_getBrokerDynamicConfig() throws Exception { final BrokerDynamicConfig brokerDynamicConfig = BrokerDynamicConfig.builder().withQueryBlocklist( List.of( - new QueryBlocklistRule("test", Set.of("dataSource"), null, null) + new DefaultQueryBlocklistRule("test", Set.of("dataSource"), null, null) ) ).build(); diff --git a/server/src/test/java/org/apache/druid/server/QueryBlocklistRuleTest.java b/server/src/test/java/org/apache/druid/server/QueryBlocklistRuleTest.java index 1594143b938c..21c2b3852a4b 100644 --- a/server/src/test/java/org/apache/druid/server/QueryBlocklistRuleTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryBlocklistRuleTest.java @@ -19,10 +19,12 @@ package org.apache.druid.server; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import org.apache.druid.query.Druids; import org.apache.druid.query.timeseries.TimeseriesQuery; +import org.apache.druid.segment.TestHelper; import org.junit.Assert; import org.junit.Test; import java.util.Map; @@ -36,7 +38,7 @@ public void testMatchAllCriteria_rejectsNullCriteria() // Rule with all null criteria would block ALL queries - this should be rejected Assert.assertThrows( IllegalArgumentException.class, - () -> new QueryBlocklistRule("match-all", null, null, null) + () -> new DefaultQueryBlocklistRule("match-all", null, null, null) ); } @@ -46,7 +48,7 @@ public void testMatchAllCriteria_rejectsEmptyCollections() // Rule with all empty collections should also be rejected (same as null) Assert.assertThrows( IllegalArgumentException.class, - () -> new QueryBlocklistRule("match-all", ImmutableSet.of(), ImmutableSet.of(), ImmutableMap.of()) + () -> new DefaultQueryBlocklistRule("match-all", ImmutableSet.of(), ImmutableSet.of(), ImmutableMap.of()) ); } @@ -54,7 +56,7 @@ public void testMatchAllCriteria_rejectsEmptyCollections() public void testMatchByDataSource() { Set dataSources = ImmutableSet.of("sensitive_data", "pii_table"); - QueryBlocklistRule rule = new QueryBlocklistRule("block-sensitive", dataSources, null, null); + QueryBlocklistRule rule = new DefaultQueryBlocklistRule("block-sensitive", dataSources, null, null); // Should match when datasource is in the list TimeseriesQuery matchingQuery = Druids.newTimeseriesQueryBuilder() @@ -75,7 +77,7 @@ public void testMatchByDataSource() public void testMatchByContext() { Map contextMatches = ImmutableMap.of("priority", "0", "application", "rogue-app"); - QueryBlocklistRule rule = new QueryBlocklistRule("block-rogue-app", null, null, contextMatches); + QueryBlocklistRule rule = new DefaultQueryBlocklistRule("block-rogue-app", null, null, contextMatches); // Should match when all context values match TimeseriesQuery matchingQuery = Druids.newTimeseriesQueryBuilder() @@ -105,7 +107,7 @@ public void testMatchByContext() public void testMatchByQueryType() { Set queryTypes = ImmutableSet.of("timeseries", "groupBy"); - QueryBlocklistRule rule = new QueryBlocklistRule("block-timeseries-groupby", null, queryTypes, null); + QueryBlocklistRule rule = new DefaultQueryBlocklistRule("block-timeseries-groupby", null, queryTypes, null); // Should match when query type is in the list (timeseries) TimeseriesQuery matchingQuery = Druids.newTimeseriesQueryBuilder() @@ -121,7 +123,7 @@ public void testMatchByMultipleCriteria() // Rule with multiple criteria - all must match (AND logic) Set dataSources = ImmutableSet.of("large_table"); Map contextMatches = ImmutableMap.of("priority", "0"); - QueryBlocklistRule rule = new QueryBlocklistRule( + QueryBlocklistRule rule = new DefaultQueryBlocklistRule( "block-low-priority-large-table", dataSources, null, @@ -156,7 +158,7 @@ public void testMatchByMultipleCriteria() @Test public void testWildcardBehavior_nullQueryTypes() { - QueryBlocklistRule rule = new QueryBlocklistRule( + QueryBlocklistRule rule = new DefaultQueryBlocklistRule( "block-datasource-all-types", ImmutableSet.of("blocked_ds"), null, // null means match all query types @@ -177,7 +179,7 @@ public void testRuleNameValidation_null() // Rule name cannot be null Assert.assertThrows( IllegalArgumentException.class, - () -> new QueryBlocklistRule(null, ImmutableSet.of("ds"), null, null) + () -> new DefaultQueryBlocklistRule(null, ImmutableSet.of("ds"), null, null) ); } @@ -187,7 +189,39 @@ public void testRuleNameValidation_empty() // Rule name cannot be empty Assert.assertThrows( IllegalArgumentException.class, - () -> new QueryBlocklistRule("", ImmutableSet.of("ds"), null, null) + () -> new DefaultQueryBlocklistRule("", ImmutableSet.of("ds"), null, null) ); } + + @Test + public void testDeserialize_missingType_usesDefault() throws Exception + { + ObjectMapper mapper = TestHelper.makeJsonMapper(); + String json = "{\"ruleName\":\"block-ds\",\"dataSources\":[\"foo\"]}"; + QueryBlocklistRule rule = mapper.readValue(json, QueryBlocklistRule.class); + Assert.assertTrue(rule instanceof DefaultQueryBlocklistRule); + Assert.assertEquals("block-ds", rule.getRuleName()); + } + + @Test + public void testDeserialize_explicitDefaultType() throws Exception + { + ObjectMapper mapper = TestHelper.makeJsonMapper(); + String json = "{\"type\":\"default\",\"ruleName\":\"block-ds\",\"dataSources\":[\"foo\"]}"; + QueryBlocklistRule rule = mapper.readValue(json, QueryBlocklistRule.class); + Assert.assertTrue(rule instanceof DefaultQueryBlocklistRule); + Assert.assertEquals("block-ds", rule.getRuleName()); + } + + @Test + public void testDeserialize_unrecognizedType_fails() + { + ObjectMapper mapper = TestHelper.makeJsonMapper(); + String json = "{\"type\":\"customExtension\",\"ruleName\":\"block-ds\",\"dataSources\":[\"foo\"]}"; + Exception e = Assert.assertThrows( + Exception.class, + () -> mapper.readValue(json, QueryBlocklistRule.class) + ); + Assert.assertTrue(e.getMessage().contains("Could not resolve type id 'customExtension'")); + } } diff --git a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java index 96880147f40b..5f8da0d21a39 100644 --- a/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java +++ b/server/src/test/java/org/apache/druid/server/QueryLifecycleTest.java @@ -830,7 +830,7 @@ public void appendTo(StringBuffer buffer) public void testRunSimple_queryBlocklisted() { // Create a blocklist rule that matches our test query - QueryBlocklistRule rule = new QueryBlocklistRule( + QueryBlocklistRule rule = new DefaultQueryBlocklistRule( "test-rule", ImmutableSet.of(DATASOURCE), null, @@ -880,7 +880,7 @@ public void testRunSimple_queryBlocklisted() public void testRunSimple_queryNotBlocklisted() { // Create a blocklist rule that does NOT match our test query - QueryBlocklistRule rule = new QueryBlocklistRule( + QueryBlocklistRule rule = new DefaultQueryBlocklistRule( "test-rule", ImmutableSet.of("other_datasource"), null, diff --git a/server/src/test/java/org/apache/druid/server/broker/BrokerDynamicConfigTest.java b/server/src/test/java/org/apache/druid/server/broker/BrokerDynamicConfigTest.java index 425ccbdd14ee..6a466bcd269e 100644 --- a/server/src/test/java/org/apache/druid/server/broker/BrokerDynamicConfigTest.java +++ b/server/src/test/java/org/apache/druid/server/broker/BrokerDynamicConfigTest.java @@ -26,6 +26,7 @@ import nl.jqno.equalsverifier.EqualsVerifier; import org.apache.druid.query.QueryContext; import org.apache.druid.segment.TestHelper; +import org.apache.druid.server.DefaultQueryBlocklistRule; import org.apache.druid.server.QueryBlocklistRule; import org.junit.Assert; import org.junit.Test; @@ -60,12 +61,35 @@ public void testSerde() throws Exception ); List expectedBlocklist = ImmutableList.of( - new QueryBlocklistRule("block-wikipedia", ImmutableSet.of("wikipedia"), null, null) + new DefaultQueryBlocklistRule("block-wikipedia", ImmutableSet.of("wikipedia"), null, null) ); Assert.assertEquals(expectedBlocklist, actual.getQueryBlocklist()); } + @Test + public void testSerdeWithExplicitDefaultType() throws Exception + { + String jsonStr = "{\n" + + " \"queryBlocklist\": [\n" + + " {\n" + + " \"type\": \"default\",\n" + + " \"ruleName\": \"block-wikipedia\",\n" + + " \"dataSources\": [\"wikipedia\"]\n" + + " }\n" + + " ]\n" + + "}\n"; + + BrokerDynamicConfig actual = mapper.readValue(jsonStr, BrokerDynamicConfig.class); + + Assert.assertEquals(1, actual.getQueryBlocklist().size()); + Assert.assertTrue(actual.getQueryBlocklist().get(0) instanceof DefaultQueryBlocklistRule); + Assert.assertEquals( + new DefaultQueryBlocklistRule("block-wikipedia", ImmutableSet.of("wikipedia"), null, null), + actual.getQueryBlocklist().get(0) + ); + } + @Test public void testSerdeWithNullBlocklist() throws Exception { @@ -108,11 +132,11 @@ public void testSerdeWithComplexBlocklist() throws Exception Assert.assertNotNull(actual.getQueryBlocklist()); Assert.assertEquals(2, actual.getQueryBlocklist().size()); - QueryBlocklistRule rule1 = actual.getQueryBlocklist().get(0); + DefaultQueryBlocklistRule rule1 = (DefaultQueryBlocklistRule) actual.getQueryBlocklist().get(0); Assert.assertEquals("block-scan-queries", rule1.getRuleName()); Assert.assertEquals(ImmutableSet.of("scan"), rule1.getQueryTypes()); - QueryBlocklistRule rule2 = actual.getQueryBlocklist().get(1); + DefaultQueryBlocklistRule rule2 = (DefaultQueryBlocklistRule) actual.getQueryBlocklist().get(1); Assert.assertEquals("block-context", rule2.getRuleName()); Assert.assertEquals(ImmutableMap.of("priority", "0"), rule2.getContextMatches()); }