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

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

paul-rogers commented on code in PR #13071:
URL: https://github.com/apache/druid/pull/13071#discussion_r991586283


##########
sql/src/main/java/org/apache/druid/sql/AbstractStatement.java:
##########
@@ -49,6 +57,14 @@ public abstract class AbstractStatement implements Closeable
   protected final SqlToolbox sqlToolbox;
   protected final SqlQueryPlus queryPlus;
   protected final SqlExecutionReporter reporter;
+
+  /**
+   * Copy of the query context modified during planning. Modifications are
+   * valid in tasks that run in the request thread. Once the query forks
+   * child threads, then concurrent modifications to the query context will
+   * result in an undefined muddle of race conditions.
+   */
+  protected final Map<String, Object> queryContext;

Review Comment:
   Yes. Revised the comment to make the semantics a bit clearer. See if the revised text helps.



##########
processing/src/main/java/org/apache/druid/query/QueryContext.java:
##########
@@ -29,227 +33,545 @@
 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.
+ * Code builds up a map of context values from serialization or during
+ * planning. Once that map is handed to the {@code QueryContext}, that map
+ * is effectively immutable.
+ * <p>
+ * The implementation uses a {@link TreeMap} so that the serialized form of a query
+ * lists context values in a deterministic order. Jackson will call
+ * {@code getContext()} on the query, which will call {@link #asMap()} here,
+ * which returns the sorted {@code TreeMap}.
+ * <p>
+ * The {@code TreeMap} is a mutable class. We'd prefer an immutable class, but
+ * we can choose either ordering or immutability. Since the semantics of the context
+ * is that it is immutable once it is placed in a query. Code should NEVER get the
+ * context map from a query and modify it, even if the actual implementation
+ * allows it.
  */
 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)
+  {
+    // There is no semantic difference between an empty and a null context.
+    // Ensure that a context always exists to avoid the need to check for
+    // a null context. Jackson serialization will omit empty contexts.
+    this.context = context == null ? Collections.emptyMap() : new TreeMap<>(context);

Review Comment:
   Thanks for this suggestion. I verified that `Collections.unmodifiableMap(.)` creates a wrapper around the `TreeMap`, which should preserve the key order which is needed for deterministic serialization.



##########
sql/src/main/java/org/apache/druid/sql/avatica/DruidConnection.java:
##########
@@ -97,7 +96,7 @@ public DruidJdbcStatement createStatement(SqlStatementFactory sqlStatementFactor
       final DruidJdbcStatement statement = new DruidJdbcStatement(
           connectionId,
           statementId,
-          context.copy(),
+          new HashMap<String, Object>(context),

Review Comment:
   Thanks for calling this to my attention. I did a scrub of the JDBC code to streamline the way we handle contexts. Please take a look to see if I missed anything.



##########
sql/src/main/java/org/apache/druid/sql/SqlQueryPlus.java:
##########
@@ -39,25 +39,31 @@
  * SQL requests come from a variety of sources in a variety of formats. Use
  * the {@link Builder} class to create an instance from the information
  * available at each point in the code.
+ * <p>
+ * The query context has a complex lifecycle. The copy here should remain
+ * unchanged: this is the set of values which the user requested. Planning will
+ * add (and sometimes remove) values: that work should be done on a copy of the
+ * context so that we have a clean record of the user's original requested
+ * values.
  */
 public class SqlQueryPlus
 {
   private final String sql;
-  private final QueryContext queryContext;
+  private final Map<String, Object> queryContext;

Review Comment:
   Clarified that the context here is immutable. Used your `Collections.unmodifableMap(.)` trick to make it so. Doing this eliminates one more area of ambiguity: it is up to the planner to make a copy before it starts modifying the context. Explained that this rule is needed so that we do security checks on the user's original keys, but not on planner-inserted keys.



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