[GitHub] [druid] FrankChen021 commented on a diff in pull request #13071: Redesign QueryContext class

[GitBox <gi...@apache.org>]

FrankChen021 commented on code in PR #13071:
URL: https://github.com/apache/druid/pull/13071#discussion_r985034097


##########
processing/src/main/java/org/apache/druid/query/scan/ScanQueryEngine.java:
##########
@@ -46,6 +45,7 @@
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
+

Review Comment:
   ```suggestion
   
   ```



##########
server/src/main/java/org/apache/druid/server/QueryResource.java:
##########
@@ -386,14 +386,7 @@ private Query<?> readQuery(
     String prevEtag = getPreviousEtag(req);
 
     if (prevEtag != null) {
-      if (baseQuery.getQueryContext() == null) {
-        QueryContext context = new QueryContext(baseQuery.getContext());
-        context.addSystemParam(HEADER_IF_NONE_MATCH, prevEtag);
-
-        return baseQuery.withOverriddenContext(context.getMergedParams());
-      } else {
-        baseQuery.getQueryContext().addSystemParam(HEADER_IF_NONE_MATCH, prevEtag);
-      }
+      baseQuery.getContext().put(HEADER_IF_NONE_MATCH, prevEtag);

Review Comment:
   `getContext()` returns an immutable map according to current implementation, a call of `put` will result in exception



##########
processing/src/main/java/org/apache/druid/query/select/SelectQuery.java:
##########
@@ -34,6 +33,7 @@
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
+

Review Comment:
   ```suggestion
   
   ```



##########
processing/src/main/java/org/apache/druid/query/search/SearchQueryQueryToolChest.java:
##########
@@ -46,6 +45,7 @@
 import org.apache.druid.query.dimension.DimensionSpec;
 
 import javax.annotation.Nullable;
+

Review Comment:
   ```suggestion
   
   ```



##########
processing/src/main/java/org/apache/druid/segment/VirtualColumns.java:
##########
@@ -48,6 +47,7 @@
 import org.apache.druid.segment.virtual.VirtualizedColumnSelectorFactory;
 
 import javax.annotation.Nullable;
+

Review Comment:
   ```suggestion
   
   ```



##########
server/src/main/java/org/apache/druid/client/DirectDruidClient.java:
##########
@@ -152,21 +153,23 @@ public Sequence<T> run(final QueryPlus<T> queryPlus, final ResponseContext conte
   {
     final Query<T> query = queryPlus.getQuery();
     QueryToolChest<T, Query<T>> toolChest = warehouse.getToolChest(query);
-    boolean isBySegment = QueryContexts.isBySegment(query);
+    boolean isBySegment = query.context().isBySegment();
     final JavaType queryResultType = isBySegment ? toolChest.getBySegmentResultType() : toolChest.getBaseResultType();
 
     final ListenableFuture<InputStream> future;
     final String url = scheme + "://" + host + "/druid/v2/";
     final String cancelUrl = url + query.getId();
 
     try {
-      log.debug("Querying queryId[%s] url[%s]", query.getId(), url);
+      log.debug("Querying queryId [%s] url [%s]", query.getId(), url);
 
       final long requestStartTimeNs = System.nanoTime();
-      final long timeoutAt = query.getQueryContext().getAsLong(QUERY_FAIL_TIME);
-      final long maxScatterGatherBytes = QueryContexts.getMaxScatterGatherBytes(query);
+      final QueryContext queryContext = query.context();
+      // Will NPE if the value is not set.

Review Comment:
   What does this mean? Where will the NPE be thrown?



##########
server/src/main/java/org/apache/druid/server/ClientQuerySegmentWalker.java:
##########
@@ -60,6 +59,7 @@
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
+

Review Comment:
   ```suggestion
   
   ```



##########
server/src/main/java/org/apache/druid/server/QueryResource.java:
##########
@@ -78,6 +77,7 @@
 import javax.ws.rs.core.Response;
 import javax.ws.rs.core.Response.Status;
 import javax.ws.rs.core.StreamingOutput;
+

Review Comment:
   ```suggestion
   
   ```



##########
processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQueryEngine.java:
##########
@@ -49,6 +48,7 @@
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
+

Review Comment:
   ```suggestion
   
   ```



##########
processing/src/main/java/org/apache/druid/query/QueryContext.java:
##########
@@ -20,236 +20,540 @@
 package org.apache.druid.query;
 
 import org.apache.druid.java.util.common.HumanReadableBytes;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.granularity.Granularity;
+import org.apache.druid.query.QueryContexts.Vectorize;
+import org.apache.druid.segment.QueryableIndexStorageAdapter;
 
 import javax.annotation.Nullable;
 
 import java.util.Collections;
 import java.util.Map;
 import java.util.Objects;
-import java.util.TreeMap;
 
 /**
- * Holder for query context parameters. There are 3 ways to set context params today.
- *
- * - Default parameters. These are set mostly via {@link DefaultQueryConfig#context}.
- *   Auto-generated queryId or sqlQueryId are also set as default parameters. These default parameters can
- *   be overridden by user or system parameters.
- * - User parameters. These are the params set by the user. User params override default parameters but
- *   are overridden by system parameters.
- * - System parameters. These are the params set by the Druid query engine for internal use only.
- *
- * You can use {@code getX} methods or {@link #getMergedParams()} to compute the context params
- * merging 3 types of params above.
- *
- * Currently, this class is mainly used for query context parameter authorization,
- * such as HTTP query endpoints or JDBC endpoint. Its usage can be expanded in the future if we
- * want to track user parameters and separate them from others during query processing.
+ * Immutable holder for query context parameters with typed access methods.
  */
 public class QueryContext
 {
-  private final Map<String, Object> defaultParams;
-  private final Map<String, Object> userParams;
-  private final Map<String, Object> systemParams;
+  private static final QueryContext EMPTY = new QueryContext(null);
+
+  private final Map<String, Object> context;
+
+  public QueryContext(Map<String, Object> context)
+  {
+    this.context = context == null ? Collections.emptyMap() : Collections.unmodifiableMap(context);
+  }
+
+  public static QueryContext empty()
+  {
+    return EMPTY;
+  }
+
+  public static QueryContext of(Map<String, Object> context)
+  {
+    return new QueryContext(context == null ? Collections.emptyMap() : context);
+  }
+
+  public boolean isEmpty()
+  {
+    return context.isEmpty();
+  }
+
+  public Map<String, Object> getContext()
+  {
+    return context;
+  }
 
   /**
-   * Cache of params merged.
+   * Check if the given key is set. If the client will then fetch the value,
+   * consider using one of the {@code get<Type>(String key)} methods instead:
+   * they each return {@code null} if the value is not set.
+   */
+  public boolean containsKey(String key)
+  {
+    return context.containsKey(key);
+  }
+
+  /**
+   * Return a value as a generic {@code Object}, returning {@code null} if the
+   * context value is not set.
    */
   @Nullable
-  private Map<String, Object> mergedParams;
+  public Object get(String key)
+  {
+    return context.get(key);
+  }
 
-  public QueryContext()
+  /**
+   * Return a value as a generic {@code Object}, returning the default value if the
+   * context value is not set.
+   */
+  public Object get(String key, Object defaultValue)
   {
-    this(null);
+    final Object val = get(key);
+    return val == null ? defaultValue : val;
   }
 
-  public QueryContext(@Nullable Map<String, Object> userParams)
+  /**
+   * Return a value as an {@code String}, returning {@link null} if the
+   * context value is not set.
+   *
+   * @throws BadQueryContextException for an invalid value
+   */
+  @Nullable
+  public String getString(String key)
   {
-    this(
-        new TreeMap<>(),
-        userParams == null ? new TreeMap<>() : new TreeMap<>(userParams),
-        new TreeMap<>()
-    );
+    return getString(key, null);
   }
 
-  private QueryContext(
-      final Map<String, Object> defaultParams,
-      final Map<String, Object> userParams,
-      final Map<String, Object> systemParams
-  )
+  public String getString(String key, String defaultValue)
   {
-    this.defaultParams = defaultParams;
-    this.userParams = userParams;
-    this.systemParams = systemParams;
-    this.mergedParams = null;
+    return QueryContexts.parseString(context, key, defaultValue);
   }
 
-  private void invalidateMergedParams()
+  /**
+   * Return a value as an {@code Boolean}, returning {@link null} if the
+   * context value is not set.
+   *
+   * @throws BadQueryContextException for an invalid value
+   */
+  public Boolean getBoolean(final String key)
   {
-    this.mergedParams = null;
+    return QueryContexts.getAsBoolean(key, get(key));
   }
 
-  public boolean isEmpty()
+  /**
+   * Return a value as an {@code boolean}, returning the default value if the
+   * context value is not set.
+   *
+   * @throws BadQueryContextException for an invalid value
+   */
+  public boolean getBoolean(final String key, final boolean defaultValue)
   {
-    return defaultParams.isEmpty() && userParams.isEmpty() && systemParams.isEmpty();
+    return QueryContexts.parseBoolean(context, key, defaultValue);
   }
 
-  public void addDefaultParam(String key, Object val)
+  /**
+   * Return a value as an {@code Integer}, returning {@link null} if the
+   * context value is not set.
+   *
+   * @throws BadQueryContextException for an invalid value
+   */
+  public Integer getInt(final String key)
+  {
+    return QueryContexts.getAsInt(key, get(key));
+  }
+
+  /**
+   * Return a value as an {@code int}, returning the default value if the
+   * context value is not set.
+   *
+   * @throws BadQueryContextException for an invalid value
+   */
+  public int getInt(final String key, final int defaultValue)
   {
-    invalidateMergedParams();
-    defaultParams.put(key, val);
+    return QueryContexts.parseInt(context, key, defaultValue);
   }
 
-  public void addDefaultParams(Map<String, Object> defaultParams)
+  /**
+   * Return a value as an {@code Long}, returning {@link null} if the
+   * context value is not set.
+   *
+   * @throws BadQueryContextException for an invalid value
+   */
+  public Long getLong(final String key)
   {
-    invalidateMergedParams();
-    this.defaultParams.putAll(defaultParams);
+    return QueryContexts.getAsLong(key, get(key));
   }
 
-  public void addSystemParam(String key, Object val)
+  /**
+   * Return a value as an {@code long}, returning the default value if the
+   * context value is not set.
+   *
+   * @throws BadQueryContextException for an invalid value
+   */
+  public long getLong(final String key, final long defaultValue)
   {
-    invalidateMergedParams();
-    this.systemParams.put(key, val);
+    return QueryContexts.parseLong(context, key, defaultValue);
   }
 
-  public Object removeUserParam(String key)
+  /**
+   * Return a value as an {@code Float}, returning {@link null} if the
+   * context value is not set.
+   *
+   * @throws BadQueryContextException for an invalid value
+   */
+  @SuppressWarnings("unused")
+  public Float getFloat(final String key)
   {
-    invalidateMergedParams();
-    return userParams.remove(key);
+    return QueryContexts.getAsFloat(key, get(key));
   }
 
   /**
-   * Returns only the context parameters the user sets.
-   * The returned map does not include the parameters that have been removed via {@link #removeUserParam}.
+   * Return a value as an {@code long}, returning the default value if the

Review Comment:
   ```suggestion
      * Return a value as an {@code float}, returning the default value if the
   ```



##########
server/src/main/java/org/apache/druid/server/QueryLifecycle.java:
##########
@@ -61,7 +62,9 @@
 
 import javax.annotation.Nullable;
 import javax.servlet.http.HttpServletRequest;
+

Review Comment:
   ```suggestion
   
   ```



##########
sql/src/main/java/org/apache/druid/sql/calcite/planner/QueryHandler.java:
##########
@@ -537,8 +536,7 @@ protected PlannerResult planWithDruidConvention() throws ValidationException
   @Nullable
   private RelRoot possiblyWrapRootWithOuterLimitFromContext(RelRoot root)
   {
-    Object outerLimitObj = handlerContext.queryContext().get(PlannerContext.CTX_SQL_OUTER_LIMIT);
-    Long outerLimit = DimensionHandlerUtils.convertObjectToLong(outerLimitObj, true);
+    Long outerLimit = handlerContext.queryContext().getLong(PlannerContext.CTX_SQL_OUTER_LIMIT);

Review Comment:
   Original code calls `DimensionHandlerUtils.convertObjectToLong` to convert an object to Long. When the input is type of String, `DimensionHandlerUtils.getExactLongFromDecimalString` is called to parse the string. This differs from `queryContext().getLong` which calls `Numbers.parseLong`. 
   
   I'm not sure if there's parse difference between these two parse method.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org