Avoid ClassCastException when getting values from QueryContext

extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQuery.java

@@ -24,7 +24,6 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Ordering;
-import org.apache.druid.java.util.common.HumanReadableBytes;
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.query.BaseQuery;
 import org.apache.druid.query.DataSource;
@@ -47,20 +46,20 @@
 import java.util.Objects;
 
 /**
- * MaterializedViewQuery helps to do materialized view selection automatically. 
- * 
+ * MaterializedViewQuery helps to do materialized view selection automatically.
+ *
  * Each MaterializedViewQuery contains a real query which type can be topn, timeseries or groupBy.
  * The real query will be optimized based on its dataSources and intervals. It will be converted into one or more
  * sub-queries, in which dataSources and intervals are replaced by derived dataSources and related sub-intervals.
- * 
+ *
  * Derived dataSources always have less dimensions, but contains all dimensions which real query required.
  */
-public class MaterializedViewQuery<T> implements Query<T> 
+public class MaterializedViewQuery<T> implements Query<T>
 {
   public static final String TYPE = "view";
   private final Query query;
   private final DataSourceOptimizer optimizer;
-  
+
   @JsonCreator
   public MaterializedViewQuery(
       @JsonProperty("query") Query query,
@@ -74,24 +73,24 @@ public MaterializedViewQuery(
     this.query = query;
     this.optimizer = optimizer;
   }
-  
+
   @JsonProperty("query")
   public Query getQuery()
   {
     return query;
   }
-  
+
   public DataSourceOptimizer getOptimizer()
   {
     return optimizer;
   }
-  
+
   @Override
   public DataSource getDataSource()
   {
     return query.getDataSource();
   }
-  
+
   @Override
   public boolean hasFilters()
   {
@@ -111,14 +110,14 @@ public String getType()
   }
 
   @Override
-  public QueryRunner<T> getRunner(QuerySegmentWalker walker) 
+  public QueryRunner<T> getRunner(QuerySegmentWalker walker)
   {
     return ((BaseQuery) query).getQuerySegmentSpec().lookup(this, walker);
   }
 
   @Override
   public List<Interval> getIntervals()
-      
+
   {
     return query.getIntervals();
   }
@@ -153,30 +152,6 @@ public QueryContext getQueryContext()
     return query.getQueryContext();
   }
 
-  @Override
-  public <ContextType> ContextType getContextValue(String key)
-  {
-    return (ContextType) query.getContextValue(key);
-  }
-
-  @Override
-  public <ContextType> ContextType getContextValue(String key, ContextType defaultValue) 
-  {
-    return (ContextType) query.getContextValue(key, defaultValue);
-  }
-
-  @Override
-  public boolean getContextBoolean(String key, boolean defaultValue) 
-  {
-    return query.getContextBoolean(key, defaultValue);
-  }
-
-  @Override
-  public HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue)
-  {
-    return query.getContextHumanReadableBytes(key, defaultValue);
-  }
-
   @Override
   public boolean isDescending()
   {
@@ -190,13 +165,13 @@ public Ordering<T> getResultOrdering()
   }
 
   @Override
-  public MaterializedViewQuery withOverriddenContext(Map<String, Object> contextOverride) 
+  public MaterializedViewQuery withOverriddenContext(Map<String, Object> contextOverride)
   {
     return new MaterializedViewQuery(query.withOverriddenContext(contextOverride), optimizer);
   }
 
   @Override
-  public MaterializedViewQuery withQuerySegmentSpec(QuerySegmentSpec spec) 
+  public MaterializedViewQuery withQuerySegmentSpec(QuerySegmentSpec spec)
   {
     return new MaterializedViewQuery(query.withQuerySegmentSpec(spec), optimizer);
   }
@@ -227,7 +202,7 @@ public String getSubQueryId()
   }
 
   @Override
-  public MaterializedViewQuery withDataSource(DataSource dataSource) 
+  public MaterializedViewQuery withDataSource(DataSource dataSource)
   {
     return new MaterializedViewQuery(query.withDataSource(dataSource), optimizer);
   }

extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryTest.java

@@ -121,7 +121,7 @@ public void testGetContextHumanReadableBytes()
         .postAggregators(QueryRunnerTestHelper.ADD_ROWS_INDEX_CONSTANT)
         .build();
     MaterializedViewQuery query = new MaterializedViewQuery(topNQuery, optimizer);
-    Assert.assertEquals(20_000_000, query.getContextHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).getBytes());
+    Assert.assertEquals(20_000_000, query.getContextAsHumanReadableBytes("maxOnDiskStorage", HumanReadableBytes.ZERO).getBytes());
 
   }
 }

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java

@@ -1454,7 +1454,7 @@ private static GranularitySpec makeGranularitySpecForIngestion(
   )
   {
     if (isRollupQuery) {
-      final String queryGranularity = query.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, "");
+      final String queryGranularity = query.getQueryContext().getAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, "");
 
       if (timeIsGroupByDimension((GroupByQuery) query, columnMappings) && !queryGranularity.isEmpty()) {
         return new ArbitraryGranularitySpec(
@@ -1483,7 +1483,7 @@ private static boolean timeIsGroupByDimension(GroupByQuery groupByQuery, ColumnM
   {
     if (columnMappings.hasOutputColumn(ColumnHolder.TIME_COLUMN_NAME)) {
       final String queryTimeColumn = columnMappings.getQueryColumnForOutputColumn(ColumnHolder.TIME_COLUMN_NAME);
-      return queryTimeColumn.equals(groupByQuery.getContextValue(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD));
+      return queryTimeColumn.equals(groupByQuery.getQueryContext().getAsString(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD));
     } else {
       return false;
     }

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/QueryKitUtils.java

@@ -191,7 +191,7 @@ public static RowSignature sortableSignature(
   public static VirtualColumn makeSegmentGranularityVirtualColumn(final Query<?> query)
   {
     final Granularity segmentGranularity = QueryKitUtils.getSegmentGranularityFromContext(query.getContext());
-    final String timeColumnName = query.getContextValue(QueryKitUtils.CTX_TIME_COLUMN_NAME);
+    final String timeColumnName = query.getQueryContext().getAsString(QueryKitUtils.CTX_TIME_COLUMN_NAME);
 
     if (timeColumnName == null || Granularities.ALL.equals(segmentGranularity)) {
       return null;

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryKit.java

@@ -57,7 +57,7 @@ public static RowSignature getAndValidateSignature(final ScanQuery scanQuery, fi
   {
     RowSignature scanSignature;
     try {
-      final String s = scanQuery.getContextValue(DruidQuery.CTX_SCAN_SIGNATURE);
+      final String s = scanQuery.getQueryContext().getAsString(DruidQuery.CTX_SCAN_SIGNATURE);
       scanSignature = jsonMapper.readValue(s, RowSignature.class);
     }
     catch (JsonProcessingException e) {

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

@@ -181,27 +181,14 @@ public QueryContext getQueryContext()
     return context;
   }
 
-  @Override
-  public <ContextType> ContextType getContextValue(String key)
-  {
-    return (ContextType) context.get(key);
-  }
-
-  @Override
-  public <ContextType> ContextType getContextValue(String key, ContextType defaultValue)
-  {
-    ContextType retVal = getContextValue(key);
-    return retVal == null ? defaultValue : retVal;
-  }
-
   @Override
   public boolean getContextBoolean(String key, boolean defaultValue)
   {
     return context.getAsBoolean(key, defaultValue);
   }
 
   @Override
-  public HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue)
+  public HumanReadableBytes getContextAsHumanReadableBytes(String key, HumanReadableBytes defaultValue)
   {
     return context.getAsHumanReadableBytes(key, defaultValue);
   }

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

@@ -124,27 +124,53 @@ default QueryContext getQueryContext()
     return null;
   }
 
-  <ContextType> ContextType getContextValue(String key);
-
-  <ContextType> ContextType getContextValue(String key, ContextType defaultValue);
+  /**
+   * Get context value and cast to ContextType in an unsafe way.
+   *
+   * For safe conversion, it's recommended to use following methods instead
+   *
+   * {@link QueryContext#getAsBoolean(String)}
+   * {@link QueryContext#getAsString(String)}
+   * {@link QueryContext#getAsInt(String)}
+   * {@link QueryContext#getAsLong(String)}
+   * {@link QueryContext#getAsFloat(String, float)}
+   * {@link QueryContext#getAsEnum(String, Class, Enum)}
+   * {@link QueryContext#getAsHumanReadableBytes(String, HumanReadableBytes)}
+   */
+  @Nullable
+  default <ContextType> ContextType getContextValue(String key)
+  {
+    if (getQueryContext() == null) {
+      return null;
+    } else {
+      return (ContextType) getQueryContext().get(key);
+    }
+  }
 
-  boolean getContextBoolean(String key, boolean defaultValue);
+  default boolean getContextBoolean(String key, boolean defaultValue)
+  {
+    if (getQueryContext() == null) {
+      return defaultValue;
+    } else {
+      return getQueryContext().getAsBoolean(key, defaultValue);
+    }
+  }
 
   /**
    * Returns {@link HumanReadableBytes} for a specified context key. If the context is null or the key doesn't exist
    * a caller specified default value is returned. A default implementation is provided since Query is an extension
    * point. Extensions can choose to rely on this default to retain compatibility with core Druid.
    *
-   * @param key The context key value being looked up
+   * @param key          The context key value being looked up
    * @param defaultValue The default to return if the key value doesn't exist or the context is null.
    * @return {@link HumanReadableBytes}
    */
-  default HumanReadableBytes getContextHumanReadableBytes(String key, HumanReadableBytes defaultValue)
+  default HumanReadableBytes getContextAsHumanReadableBytes(String key, HumanReadableBytes defaultValue)
   {
-    if (null != getQueryContext()) {
-      return getQueryContext().getAsHumanReadableBytes(key, defaultValue);
-    } else {
+    if (getQueryContext() == null) {
       return defaultValue;
+    } else {
+      return getQueryContext().getAsHumanReadableBytes(key, defaultValue);
     }
   }
 
@@ -204,7 +230,7 @@ default Query<T> withSqlQueryId(String sqlQueryId)
   @Nullable
   default String getSqlQueryId()
   {
-    return getContextValue(BaseQuery.SQL_QUERY_ID);
+    return getQueryContext().getAsString(BaseQuery.SQL_QUERY_ID);
   }
 
   /**

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

@@ -168,33 +168,66 @@ public Object getOrDefault(String key, Object defaultValue)
   @Nullable
   public String getAsString(String key)
   {
-    return (String) get(key);
+    Object val = get(key);
+    return val == null ? null : val.toString();
+  }
+
+  public String getAsString(String key, String defaultValue)
+  {
+    Object val = get(key);
+    return val == null ? defaultValue : val.toString();
+  }
+
+  @Nullable
+  public Boolean getAsBoolean(String key)
+  {
+    return QueryContexts.getAsBoolean(key, get(key));
   }
 
   public boolean getAsBoolean(
-      final String parameter,
+      final String key,
       final boolean defaultValue
   )
   {
-    return QueryContexts.getAsBoolean(parameter, get(parameter), defaultValue);
+    return QueryContexts.getAsBoolean(key, get(key), defaultValue);
+  }
+
+  public Integer getAsInt(final String key)
+  {
+    return QueryContexts.getAsInt(key, get(key));
   }
 
   public int getAsInt(
-      final String parameter,
+      final String key,
       final int defaultValue
   )
   {
-    return QueryContexts.getAsInt(parameter, get(parameter), defaultValue);
+    return QueryContexts.getAsInt(key, get(key), defaultValue);
+  }
+
+  public Long getAsLong(final String key)
+  {
+    return QueryContexts.getAsLong(key, get(key));
+  }
+
+  public long getAsLong(final String key, final long defaultValue)
+  {
+    return QueryContexts.getAsLong(key, get(key), defaultValue);
+  }
+
+  public HumanReadableBytes getAsHumanReadableBytes(final String key, final HumanReadableBytes defaultValue)
+  {
+    return QueryContexts.getAsHumanReadableBytes(key, get(key), defaultValue);
   }
 
-  public long getAsLong(final String parameter, final long defaultValue)
+  public float getAsFloat(final String key, final float defaultValue)
   {
-    return QueryContexts.getAsLong(parameter, get(parameter), defaultValue);
+    return QueryContexts.getAsFloat(key, get(key), defaultValue);
   }
 
-  public HumanReadableBytes getAsHumanReadableBytes(final String parameter, final HumanReadableBytes defaultValue)
+  public <E extends Enum<E>> E getAsEnum(String key, Class<E> clazz, E defaultValue)
   {
-    return QueryContexts.getAsHumanReadableBytes(parameter, get(parameter), defaultValue);
+    return QueryContexts.getAsEnum(key, get(key), clazz, defaultValue);
   }
 
   public Map<String, Object> getMergedParams()