Apache
/
druid
mirror of https://github.com/apache/druid.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Clean up query contexts (#12633) 

* Clean up query contexts

Uses constants in place of literal strings for context keys.
Moves some QueryContext methods to QueryContexts for reuse.

* Revisions from review comments
This commit is contained in:
Paul Rogers 2022-06-15 11:31:22 -07:00 committed by GitHub
parent 28f2c8e112
commit 45e3111549
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 209 additions and 90 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
processing/src/main/java/org/apache/druid/query/BaseQuery.java
View File

@ -35,10 +35,10 @@ import org.joda.time.Duration;
import org.joda.time.Interval;
import javax.annotation.Nullable;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.TreeMap;
/**
*
@ -210,13 +210,7 @@ public abstract class BaseQuery<T> implements Query<T>
final Map<String, Object> overrides
)
{
Map<String, Object> overridden = new TreeMap<>();
if (context != null) {
overridden.putAll(context);
}
overridden.putAll(overrides);
return overridden;
return QueryContexts.override(context, overrides);
}
/**

2
processing/src/main/java/org/apache/druid/query/FinalizeResultsQueryRunner.java
View File

@ -64,7 +64,7 @@ public class FinalizeResultsQueryRunner<T> implements QueryRunner<T>
final MetricManipulationFn metricManipulationFn;
if (shouldFinalize) {
queryToRun = query.withOverriddenContext(ImmutableMap.of("finalize", false));
queryToRun = query.withOverriddenContext(ImmutableMap.of(QueryContexts.FINALIZE_KEY, false));
metricManipulationFn = MetricManipulatorFns.finalizing();
} else {
queryToRun = query;

37
processing/src/main/java/org/apache/druid/query/QueryContext.java
View File

@ -19,10 +19,8 @@
package org.apache.druid.query;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.Numbers;
import javax.annotation.Nullable;
import java.util.Collections;
import java.util.Map;
import java.util.Objects;
@ -162,16 +160,7 @@ public class QueryContext
final boolean defaultValue
)
{
final Object value = get(parameter);
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Boolean.parseBoolean((String) value);
} else if (value instanceof Boolean) {
return (Boolean) value;
} else {
throw new IAE("Expected parameter[%s] to be boolean", parameter);
}
return QueryContexts.getAsBoolean(parameter, get(parameter), defaultValue);
}
public int getAsInt(
@ -179,30 +168,12 @@ public class QueryContext
final int defaultValue
)
{
final Object value = get(parameter);
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Numbers.parseInt(value);
} else if (value instanceof Number) {
return ((Number) value).intValue();
} else {
throw new IAE("Expected parameter[%s] to be integer", parameter);
}
return QueryContexts.getAsInt(parameter, get(parameter), defaultValue);
}
public long getAsLong(final String parameter, final long defaultValue)
{
final Object value = get(parameter);
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Numbers.parseLong(value);
} else if (value instanceof Number) {
return ((Number) value).longValue();
} else {
throw new IAE("Expected parameter[%s] to be long", parameter);
}
return QueryContexts.getAsLong(parameter, get(parameter), defaultValue);
}
public Map<String, Object> getMergedParams()

133
processing/src/main/java/org/apache/druid/query/QueryContexts.java
View File

@ -31,6 +31,7 @@ import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.segment.QueryableIndexStorageAdapter;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.TimeUnit;
@PublicApi
@ -71,6 +72,12 @@ public class QueryContexts
public static final String BROKER_SERVICE_NAME = "brokerService";
public static final String IN_SUB_QUERY_THRESHOLD_KEY = "inSubQueryThreshold";
public static final String TIME_BOUNDARY_PLANNING_KEY = "enableTimeBoundaryPlanning";
public static final String POPULATE_CACHE_KEY = "populateCache";
public static final String POPULATE_RESULT_LEVEL_CACHE_KEY = "populateResultLevelCache";
public static final String USE_RESULT_LEVEL_CACHE_KEY = "useResultLevelCache";
public static final String SERIALIZE_DATE_TIME_AS_LONG_KEY = "serializeDateTimeAsLong";
public static final String SERIALIZE_DATE_TIME_AS_LONG_INNER_KEY = "serializeDateTimeAsLongInner";
public static final String UNCOVERED_INTERVALS_LIMIT_KEY = "uncoveredIntervalsLimit";
public static final boolean DEFAULT_BY_SEGMENT = false;
public static final boolean DEFAULT_POPULATE_CACHE = true;
@ -158,7 +165,7 @@ public class QueryContexts
public static <T> boolean isPopulateCache(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "populateCache", defaultValue);
return parseBoolean(query, POPULATE_CACHE_KEY, defaultValue);
}
public static <T> boolean isUseCache(Query<T> query)
@ -178,7 +185,7 @@ public class QueryContexts
public static <T> boolean isPopulateResultLevelCache(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "populateResultLevelCache", defaultValue);
return parseBoolean(query, POPULATE_RESULT_LEVEL_CACHE_KEY, defaultValue);
}
public static <T> boolean isUseResultLevelCache(Query<T> query)
@ -188,22 +195,22 @@ public class QueryContexts
public static <T> boolean isUseResultLevelCache(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "useResultLevelCache", defaultValue);
return parseBoolean(query, USE_RESULT_LEVEL_CACHE_KEY, defaultValue);
}
public static <T> boolean isFinalize(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, FINALIZE_KEY, defaultValue);
}
public static <T> boolean isSerializeDateTimeAsLong(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "serializeDateTimeAsLong", defaultValue);
return parseBoolean(query, SERIALIZE_DATE_TIME_AS_LONG_KEY, defaultValue);
}
public static <T> boolean isSerializeDateTimeAsLongInner(Query<T> query, boolean defaultValue)
{
return parseBoolean(query, "serializeDateTimeAsLongInner", defaultValue);
return parseBoolean(query, SERIALIZE_DATE_TIME_AS_LONG_INNER_KEY, defaultValue);
}
public static <T> Vectorize getVectorize(Query<T> query)
@ -248,7 +255,7 @@ public class QueryContexts
public static <T> int getUncoveredIntervalsLimit(Query<T> query, int defaultValue)
{
return parseInt(query, "uncoveredIntervalsLimit", defaultValue);
return parseInt(query, UNCOVERED_INTERVALS_LIMIT_KEY, defaultValue);
}
public static <T> int getPriority(Query<T> query)
@ -450,32 +457,124 @@ public class QueryContexts
static <T> long parseLong(Query<T> query, String key, long defaultValue)
{
final Object val = query.getContextValue(key);
return val == null ? defaultValue : Numbers.parseLong(val);
return getAsLong(key, query.getContextValue(key), defaultValue);
}
@SuppressWarnings("unused")
static <T> long parseLong(Map<String, Object> context, String key, long defaultValue)
{
return getAsLong(key, context.get(key), defaultValue);
}
static <T> int parseInt(Query<T> query, String key, int defaultValue)
{
final Object val = query.getContextValue(key);
return val == null ? defaultValue : Numbers.parseInt(val);
return getAsInt(key, query.getContextValue(key), defaultValue);
}
static int parseInt(Map<String, Object> context, String key, int defaultValue)
{
final Object val = context.get(key);
return val == null ? defaultValue : Numbers.parseInt(val);
return getAsInt(key, context.get(key), defaultValue);
}
static <T> boolean parseBoolean(Query<T> query, String key, boolean defaultValue)
{
final Object val = query.getContextValue(key);
return val == null ? defaultValue : Numbers.parseBoolean(val);
return getAsBoolean(key, query.getContextValue(key), defaultValue);
}
static boolean parseBoolean(Map<String, Object> context, String key, boolean defaultValue)
{
final Object val = context.get(key);
return val == null ? defaultValue : Numbers.parseBoolean(val);
return getAsBoolean(key, context.get(key), defaultValue);
}
public static String getAsString(
final String parameter,
final Object value,
final String defaultValue
)
{
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return (String) value;
} else {
throw new IAE("Expected parameter [%s] to be String", parameter);
}
}
/**
* Get the value of a parameter as a {@code boolean}. The parameter is expected
* to be {@code null}, a string or a {@code Boolean} object.
*/
public static boolean getAsBoolean(
final String parameter,
final Object value,
final boolean defaultValue
)
{
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Boolean.parseBoolean((String) value);
} else if (value instanceof Boolean) {
return (Boolean) value;
} else {
throw new IAE("Expected parameter [%s] to be a boolean", parameter);
}
}
/**
* Get the value of a parameter as an {@code int}. The parameter is expected
* to be {@code null}, a string or a {@code Number} object.
*/
public static int getAsInt(
final String parameter,
final Object value,
final int defaultValue
)
{
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Numbers.parseInt(value);
} else if (value instanceof Number) {
return ((Number) value).intValue();
} else {
throw new IAE("Expected parameter [%s] to be an integer", parameter);
}
}
/**
* Get the value of a parameter as an {@code long}. The parameter is expected
* to be {@code null}, a string or a {@code Number} object.
*/
public static long getAsLong(
final String parameter,
final Object value,
final long defaultValue)
{
if (value == null) {
return defaultValue;
} else if (value instanceof String) {
return Numbers.parseLong(value);
} else if (value instanceof Number) {
return ((Number) value).longValue();
} else {
throw new IAE("Expected parameter [%s] to be a long", parameter);
}
}
public static Map<String, Object> override(
final Map<String, Object> context,
final Map<String, Object> overrides
)
{
Map<String, Object> overridden = new TreeMap<>();
if (context != null) {
overridden.putAll(context);
}
overridden.putAll(overrides);
return overridden;
}
private QueryContexts()

3
processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV1.java
View File

@ -30,6 +30,7 @@ import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.query.GroupByMergedQueryRunner;
import org.apache.druid.query.QueryContexts;
import org.apache.druid.query.QueryPlus;
import org.apache.druid.query.QueryProcessingPool;
import org.apache.druid.query.QueryRunner;
@ -115,7 +116,7 @@ public class GroupByStrategyV1 implements GroupByStrategy
.overrideContext(
ImmutableMap.<String, Object>builder()
.put(GroupByQueryConfig.CTX_KEY_STRATEGY, GroupByStrategySelector.STRATEGY_V1)
.put("finalize", false)
.put(QueryContexts.FINALIZE_KEY, false)
// Always request array result rows when passing the query down.
.put(GroupByQueryConfig.CTX_KEY_ARRAY_RESULT_ROWS, true)

2
processing/src/main/java/org/apache/druid/query/groupby/strategy/GroupByStrategyV2.java
View File

@ -214,7 +214,7 @@ public class GroupByStrategyV2 implements GroupByStrategy
// Set up downstream context.
final ImmutableMap.Builder<String, Object> context = ImmutableMap.builder();
context.put("finalize", false);
context.put(QueryContexts.FINALIZE_KEY, false);
context.put(GroupByQueryConfig.CTX_KEY_STRATEGY, GroupByStrategySelector.STRATEGY_V2);
context.put(CTX_KEY_OUTERMOST, false);

57
processing/src/test/java/org/apache/druid/query/QueryContextsTest.java
View File

@ -44,7 +44,7 @@ public class QueryContextsTest
new TableDataSource("test"),
new MultipleIntervalSegmentSpec(ImmutableList.of(Intervals.of("0/100"))),
false,
new HashMap()
new HashMap<>()
);
Assert.assertEquals(300_000, QueryContexts.getDefaultTimeout(query));
}
@ -56,7 +56,7 @@ public class QueryContextsTest
new TableDataSource("test"),
new MultipleIntervalSegmentSpec(ImmutableList.of(Intervals.of("0/100"))),
false,
new HashMap()
new HashMap<>()
);
Assert.assertEquals(300_000, QueryContexts.getTimeout(query));
@ -206,4 +206,57 @@ public class QueryContextsTest
Assert.assertTrue(QueryContexts.isDebug(query));
Assert.assertTrue(QueryContexts.isDebug(query.getContext()));
}
@Test
public void testGetAs()
{
Assert.assertNull(QueryContexts.getAsString("foo", null, null));
Assert.assertEquals("default", QueryContexts.getAsString("foo", null, "default"));
Assert.assertEquals("value", QueryContexts.getAsString("foo", "value", "default"));
try {
QueryContexts.getAsString("foo", 10, null);
Assert.fail();
}
catch (IAE e) {
// Expected
}
Assert.assertFalse(QueryContexts.getAsBoolean("foo", null, false));
Assert.assertTrue(QueryContexts.getAsBoolean("foo", null, true));
Assert.assertTrue(QueryContexts.getAsBoolean("foo", "true", false));
Assert.assertTrue(QueryContexts.getAsBoolean("foo", true, false));
try {
QueryContexts.getAsBoolean("foo", 10, false);
Assert.fail();
}
catch (IAE e) {
// Expected
}
Assert.assertEquals(10, QueryContexts.getAsInt("foo", null, 10));
Assert.assertEquals(20, QueryContexts.getAsInt("foo", "20", 10));
Assert.assertEquals(20, QueryContexts.getAsInt("foo", 20, 10));
Assert.assertEquals(20, QueryContexts.getAsInt("foo", 20L, 10));
Assert.assertEquals(20, QueryContexts.getAsInt("foo", 20D, 10));
try {
QueryContexts.getAsInt("foo", true, 20);
Assert.fail();
}
catch (IAE e) {
// Expected
}
Assert.assertEquals(10L, QueryContexts.getAsLong("foo", null, 10));
Assert.assertEquals(20L, QueryContexts.getAsLong("foo", "20", 10));
Assert.assertEquals(20L, QueryContexts.getAsLong("foo", 20, 10));
Assert.assertEquals(20L, QueryContexts.getAsLong("foo", 20L, 10));
Assert.assertEquals(20L, QueryContexts.getAsLong("foo", 20D, 10));
try {
QueryContexts.getAsLong("foo", true, 20);
Assert.fail();
}
catch (IAE e) {
// Expected
}
}
}

20
processing/src/test/java/org/apache/druid/query/datasourcemetadata/DataSourceMetadataQueryTest.java
View File

@ -79,13 +79,13 @@ public class DataSourceMetadataQueryTest
.intervals("2013/2014")
.context(
ImmutableMap.of(
"priority",
QueryContexts.PRIORITY_KEY,
1,
"useCache",
QueryContexts.USE_CACHE_KEY,
true,
"populateCache",
QueryContexts.POPULATE_CACHE_KEY,
"true",
"finalize",
QueryContexts.FINALIZE_KEY,
true
)
).build();
@ -103,12 +103,12 @@ public class DataSourceMetadataQueryTest
);
Assert.assertEquals((Integer) 1, serdeQuery.getContextValue(QueryContexts.PRIORITY_KEY));
Assert.assertEquals(true, serdeQuery.getContextValue("useCache"));
Assert.assertEquals("true", serdeQuery.getContextValue("populateCache"));
Assert.assertEquals(true, serdeQuery.getContextValue("finalize"));
Assert.assertEquals(true, serdeQuery.getContextBoolean("useCache", false));
Assert.assertEquals(true, serdeQuery.getContextBoolean("populateCache", false));
Assert.assertEquals(true, serdeQuery.getContextBoolean("finalize", false));
Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.USE_CACHE_KEY));
Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY));
Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.FINALIZE_KEY));
Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false));
Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.POPULATE_CACHE_KEY, false));
Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false));
}
@Test

2
processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryRunnerTest.java
View File

@ -7524,7 +7524,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
new LongLastAggregatorFactory("innerlast", "index", null)
)
.setGranularity(QueryRunnerTestHelper.DAY_GRAN)
.overrideContext(ImmutableMap.of("finalize", true))
.overrideContext(ImmutableMap.of(QueryContexts.FINALIZE_KEY, true))
.build();
GroupByQuery query = makeQueryBuilder()

28
processing/src/test/java/org/apache/druid/query/timeboundary/TimeBoundaryQueryTest.java
View File

@ -55,13 +55,13 @@ public class TimeBoundaryQueryTest
.intervals("2013/2014")
.context(
ImmutableMap.of(
"priority",
QueryContexts.PRIORITY_KEY,
1,
"useCache",
QueryContexts.USE_CACHE_KEY,
true,
"populateCache",
QueryContexts.POPULATE_CACHE_KEY,
true,
"finalize",
QueryContexts.FINALIZE_KEY,
true
)
).build();
@ -80,9 +80,9 @@ public class TimeBoundaryQueryTest
Assert.assertEquals(new Integer(1), serdeQuery.getContextValue(QueryContexts.PRIORITY_KEY));
Assert.assertEquals(true, serdeQuery.getContextValue("useCache"));
Assert.assertEquals(true, serdeQuery.getContextValue("populateCache"));
Assert.assertEquals(true, serdeQuery.getContextValue("finalize"));
Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.USE_CACHE_KEY));
Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY));
Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.FINALIZE_KEY));
}
@Test
@ -93,13 +93,13 @@ public class TimeBoundaryQueryTest
.intervals("2013/2014")
.context(
ImmutableMap.of(
"priority",
QueryContexts.PRIORITY_KEY,
"1",
"useCache",
QueryContexts.USE_CACHE_KEY,
"true",
"populateCache",
QueryContexts.POPULATE_CACHE_KEY,
"true",
"finalize",
QueryContexts.FINALIZE_KEY,
"true"
)
).build();
@ -118,8 +118,8 @@ public class TimeBoundaryQueryTest
Assert.assertEquals("1", serdeQuery.getContextValue(QueryContexts.PRIORITY_KEY));
Assert.assertEquals("true", serdeQuery.getContextValue("useCache"));
Assert.assertEquals("true", serdeQuery.getContextValue("populateCache"));
Assert.assertEquals("true", serdeQuery.getContextValue("finalize"));
Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.USE_CACHE_KEY));
Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY));
Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.FINALIZE_KEY));
}
}

5
server/src/test/java/org/apache/druid/client/CachingClusteredClientTest.java
View File

@ -96,6 +96,7 @@ import org.apache.druid.query.filter.InDimFilter;
import org.apache.druid.query.filter.OrDimFilter;
import org.apache.druid.query.filter.SelectorDimFilter;
import org.apache.druid.query.groupby.GroupByQuery;
import org.apache.druid.query.groupby.GroupByQueryConfig;
import org.apache.druid.query.groupby.ResultRow;
import org.apache.druid.query.groupby.strategy.GroupByStrategySelector;
import org.apache.druid.query.ordering.StringComparators;
@ -179,10 +180,10 @@ import java.util.stream.IntStream;
public class CachingClusteredClientTest
{
private static final ImmutableMap<String, Object> CONTEXT = ImmutableMap.of(
"finalize", false,
QueryContexts.FINALIZE_KEY, false,
// GroupBy v2 won't cache on the broker, so test with v1.
"groupByStrategy", GroupByStrategySelector.STRATEGY_V1
GroupByQueryConfig.CTX_KEY_STRATEGY, GroupByStrategySelector.STRATEGY_V1
);
private static final MultipleIntervalSegmentSpec SEG_SPEC = new MultipleIntervalSegmentSpec(ImmutableList.of());
private static final String DATA_SOURCE = "test";