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

processing/src/main/java/org/apache/druid/query/BaseQuery.java

@@ -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<>();
+    return QueryContexts.override(context, overrides);
-    if (context != null) {
-      overridden.putAll(context);
-    }
-    overridden.putAll(overrides);
-
-    return overridden;
   }
 
   /**

processing/src/main/java/org/apache/druid/query/FinalizeResultsQueryRunner.java

@@ -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;

processing/src/main/java/org/apache/druid/query/QueryContext.java

@@ -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);
+    return QueryContexts.getAsBoolean(parameter, get(parameter), 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 boolean", parameter);
-    }
   }
 
   public int getAsInt(
@@ -179,30 +168,12 @@ public class QueryContext
       final int defaultValue
   )
   {
-    final Object value = get(parameter);
+    return QueryContexts.getAsInt(parameter, get(parameter), 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 integer", parameter);
-    }
   }
 
   public long getAsLong(final String parameter, final long defaultValue)
   {
-    final Object value = get(parameter);
+    return QueryContexts.getAsLong(parameter, get(parameter), 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 long", parameter);
-    }
   }
 
   public Map<String, Object> getMergedParams()

processing/src/main/java/org/apache/druid/query/QueryContexts.java

@@ -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 getAsLong(key, query.getContextValue(key), defaultValue);
-    return val == null ? defaultValue : Numbers.parseLong(val);
+  }
+
+  @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 getAsInt(key, query.getContextValue(key), defaultValue);
-    return val == null ? defaultValue : Numbers.parseInt(val);
   }
 
   static int parseInt(Map<String, Object> context, String key, int defaultValue)
   {
-    final Object val = context.get(key);
+    return getAsInt(key, context.get(key), defaultValue);
-    return val == null ? defaultValue : Numbers.parseInt(val);
   }
 
   static <T> boolean parseBoolean(Query<T> query, String key, boolean defaultValue)
   {
-    final Object val = query.getContextValue(key);
+    return getAsBoolean(key, query.getContextValue(key), defaultValue);
-    return val == null ? defaultValue : Numbers.parseBoolean(val);
   }
 
   static boolean parseBoolean(Map<String, Object> context, String key, boolean defaultValue)
   {
-    final Object val = context.get(key);
+    return getAsBoolean(key, context.get(key), defaultValue);
-    return val == null ? defaultValue : Numbers.parseBoolean(val);
+  }
+
+  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()

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

@@ -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)

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

@@ -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);
 

processing/src/test/java/org/apache/druid/query/QueryContextsTest.java

@@ -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
+    }
+  }
 }

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

@@ -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(QueryContexts.USE_CACHE_KEY));
-    Assert.assertEquals("true", serdeQuery.getContextValue("populateCache"));
+    Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY));
-    Assert.assertEquals(true, serdeQuery.getContextValue("finalize"));
+    Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.FINALIZE_KEY));
-    Assert.assertEquals(true, serdeQuery.getContextBoolean("useCache", false));
+    Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.USE_CACHE_KEY, false));
-    Assert.assertEquals(true, serdeQuery.getContextBoolean("populateCache", false));
+    Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.POPULATE_CACHE_KEY, false));
-    Assert.assertEquals(true, serdeQuery.getContextBoolean("finalize", false));
+    Assert.assertEquals(true, serdeQuery.getContextBoolean(QueryContexts.FINALIZE_KEY, false));
   }
 
   @Test

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

@@ -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()

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

@@ -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(QueryContexts.USE_CACHE_KEY));
-    Assert.assertEquals(true, serdeQuery.getContextValue("populateCache"));
+    Assert.assertEquals(true, serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY));
-    Assert.assertEquals(true, serdeQuery.getContextValue("finalize"));
+    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(QueryContexts.USE_CACHE_KEY));
-    Assert.assertEquals("true", serdeQuery.getContextValue("populateCache"));
+    Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.POPULATE_CACHE_KEY));
-    Assert.assertEquals("true", serdeQuery.getContextValue("finalize"));
+    Assert.assertEquals("true", serdeQuery.getContextValue(QueryContexts.FINALIZE_KEY));
   }
 }

server/src/test/java/org/apache/druid/client/CachingClusteredClientTest.java

@@ -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";