Today is that day - Single pass through Calcite planner

integration-tests/src/test/java/org/apache/druid/tests/query/ITJdbcQueryTest.java

@@ -211,7 +211,7 @@ public void testJdbcPrepareStatementQuery()
     }
   }
 
-  @Test(expectedExceptions = AvaticaSqlException.class, expectedExceptionsMessageRegExp = ".* Parameter at position\\[0] is not bound")
+  @Test(expectedExceptions = AvaticaSqlException.class, expectedExceptionsMessageRegExp = ".* Parameter at position \\[0] is not bound")
   public void testJdbcPrepareStatementQueryMissingParameters() throws SQLException
   {
     for (String url : connections) {

sql/src/main/java/org/apache/druid/sql/calcite/planner/ValidationResult.java → sql/src/main/java/org/apache/calcite/prepare/BaseDruidSqlValidator.java

@@ -17,32 +17,27 @@
  * under the License.
  */
 
-package org.apache.druid.sql.calcite.planner;
+//CHECKSTYLE.OFF: PackageName - Must be in Calcite
 
-import com.google.common.collect.ImmutableSet;
-import org.apache.druid.server.security.Resource;
-import org.apache.druid.server.security.ResourceAction;
+package org.apache.calcite.prepare;
 
-import java.util.Set;
+import org.apache.calcite.adapter.java.JavaTypeFactory;
+import org.apache.calcite.sql.SqlOperatorTable;
+import org.apache.calcite.sql.validate.SqlConformance;
 
 /**
- * If an SQL query can be validated by {@link DruidPlanner}, the resulting artifact is the set of {@link Resource}
- * corresponding to the datasources and views which an authenticated request must be authorized for to process the
- * query.
+ * Extend the Java-focused {@link CalciteSqlValidator} to make it visible
+ * to Druid. {@link CalciteSqlValidator} itself is protected and can be
+ * extended only if this class is in the same Calcite package.
  */
-public class ValidationResult
+public class BaseDruidSqlValidator extends CalciteSqlValidator
 {
-  private final Set<ResourceAction> resourceActions;
-
-  public ValidationResult(
-      final Set<ResourceAction> resourceActions
-  )
-  {
-    this.resourceActions = ImmutableSet.copyOf(resourceActions);
-  }
-
-  public Set<ResourceAction> getResourceActions()
+  public BaseDruidSqlValidator(
+      SqlOperatorTable opTab,
+      CalciteCatalogReader catalogReader,
+      JavaTypeFactory typeFactory,
+      SqlConformance conformance)
   {
-    return resourceActions;
+    super(opTab, catalogReader, typeFactory, conformance);
   }
 }

sql/src/main/java/org/apache/druid/sql/SqlLifecycle.java

@@ -55,7 +55,6 @@
 import org.apache.druid.sql.calcite.planner.PlannerFactory;
 import org.apache.druid.sql.calcite.planner.PlannerResult;
 import org.apache.druid.sql.calcite.planner.PrepareResult;
-import org.apache.druid.sql.calcite.planner.ValidationResult;
 import org.apache.druid.sql.http.SqlParameter;
 import org.apache.druid.sql.http.SqlQuery;
 
@@ -69,6 +68,7 @@
 import java.util.UUID;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -106,14 +106,20 @@ public class SqlLifecycle
   @GuardedBy("stateLock")
   private State state = State.NEW;
 
-  // init during intialize
+  // init during initialize
   private String sql;
   private QueryContext queryContext;
   private List<TypedValue> parameters;
+
   // init during plan
+  /**
+   * The Druid planner follows the SQL statement through the lifecycle.
+   * The planner's state is start --> validate --> (prepare | plan).
+   */
+  private DruidPlanner planner;
   private PlannerContext plannerContext;
-  private ValidationResult validationResult;
   private PrepareResult prepareResult;
+  private Set<ResourceAction> resourceActions;
   private PlannerResult plannerResult;
 
   public SqlLifecycle(
@@ -196,14 +202,13 @@ public void validateAndAuthorize(AuthenticationResult authenticationResult)
     }
     transition(State.INITIALIZED, State.AUTHORIZING);
     validate(authenticationResult);
-    Access access = doAuthorize(
+    doAuthorize(resourceActions ->
         AuthorizationUtils.authorizeAllResourceActions(
             authenticationResult,
-            validationResult.getResourceActions(),
+            resourceActions,
             plannerFactory.getAuthorizerMapper()
         )
     );
-    checkAccess(access);
   }
 
   /**
@@ -218,26 +223,33 @@ public void validateAndAuthorize(HttpServletRequest req)
     transition(State.INITIALIZED, State.AUTHORIZING);
     AuthenticationResult authResult = AuthorizationUtils.authenticationResultFromRequest(req);
     validate(authResult);
-    Access access = doAuthorize(
+    doAuthorize(resourceActions ->
         AuthorizationUtils.authorizeAllResourceActions(
             req,
-            validationResult.getResourceActions(),
+            resourceActions,
             plannerFactory.getAuthorizerMapper()
         )
     );
-    checkAccess(access);
   }
 
-  private ValidationResult validate(AuthenticationResult authenticationResult)
+  /**
+   * Perform the validation step on the Druid planner, leaving the planner
+   * ready to perform either prepare or plan.
+   */
+  private void validate(AuthenticationResult authenticationResult)
   {
-    try (DruidPlanner planner = plannerFactory.createPlanner(sql, queryContext)) {
+    try {
+      planner = plannerFactory.createPlanner(sql, queryContext);
       // set planner context for logs/metrics in case something explodes early
-      this.plannerContext = planner.getPlannerContext();
-      this.plannerContext.setAuthenticationResult(authenticationResult);
+      plannerContext = planner.getPlannerContext();
+      plannerContext.setAuthenticationResult(authenticationResult);
       // set parameters on planner context, if parameters have already been set
-      this.plannerContext.setParameters(parameters);
-      this.validationResult = planner.validate(authConfig.authorizeQueryContextParams());
-      return validationResult;
+      plannerContext.setParameters(parameters);
+      planner.validate();
+
+      // Capture the resource actions as these are reference past the
+      // life of the planner itself.
+      resourceActions = planner.resourceActions(authConfig.authorizeQueryContextParams());
     }
     // we can't collapse catch clauses since SqlPlanningException has type-sensitive constructors.
     catch (SqlParseException e) {
@@ -248,46 +260,44 @@ private ValidationResult validate(AuthenticationResult authenticationResult)
     }
   }
 
-  private Access doAuthorize(final Access authorizationResult)
+  private void doAuthorize(Function<Set<ResourceAction>, Access> authorizer)
   {
+    Access authorizationResult = planner.authorize(
+        authorizer,
+        authConfig.authorizeQueryContextParams()
+    );
     if (!authorizationResult.isAllowed()) {
       // Not authorized; go straight to Jail, do not pass Go.
       transition(State.AUTHORIZING, State.UNAUTHORIZED);
     } else {
       transition(State.AUTHORIZING, State.AUTHORIZED);
     }
-    return authorizationResult;
-  }
-
-  private void checkAccess(Access access)
-  {
-    plannerContext.setAuthorizationResult(access);
-    if (!access.isAllowed()) {
-      throw new ForbiddenException(access.toString());
+    if (!authorizationResult.isAllowed()) {
+      throw new ForbiddenException(authorizationResult.toString());
     }
   }
 
   /**
-   * Prepare the query lifecycle for execution, without completely planning into something that is executable, but
-   * including some initial parsing and validation and any dyanmic parameter type resolution, to support prepared
+   * Prepare the query lifecycle for execution, without completely planning into
+   * something that is executable, but including some initial parsing and
+   * validation and any dynamic parameter type resolution, to support prepared
    * statements via JDBC.
+   *
+   * The planner must have already performed the validation step: the planner
+   * state is reused here.
    */
-  public PrepareResult prepare() throws RelConversionException
+  public PrepareResult prepare()
   {
     synchronized (stateLock) {
       if (state != State.AUTHORIZED) {
-        throw new ISE("Cannot prepare because current state[%s] is not [%s].", state, State.AUTHORIZED);
+        throw new ISE("Cannot prepare because current state [%s] is not [%s].", state, State.AUTHORIZED);
       }
     }
     Preconditions.checkNotNull(plannerContext, "Cannot prepare, plannerContext is null");
-    try (DruidPlanner planner = plannerFactory.createPlannerWithContext(plannerContext)) {
+    try {
       this.prepareResult = planner.prepare();
       return prepareResult;
     }
-    // we can't collapse catch clauses since SqlPlanningException has type-sensitive constructors.
-    catch (SqlParseException e) {
-      throw new SqlPlanningException(e);
-    }
     catch (ValidationException e) {
       throw new SqlPlanningException(e);
     }
@@ -296,22 +306,27 @@ public PrepareResult prepare() throws RelConversionException
   /**
    * Plan the query to enable execution.
    *
-   * If successful, the lifecycle will first transition from {@link State#AUTHORIZED} to {@link State#PLANNED}.
+   * The planner must have already performed the validation step: the planner
+   * state is reused here.
+   *
+   * If successful, the lifecycle will first transition from
+   * {@link State#AUTHORIZED} to {@link State#PLANNED}.
    */
   public void plan() throws RelConversionException
   {
     transition(State.AUTHORIZED, State.PLANNED);
     Preconditions.checkNotNull(plannerContext, "Cannot plan, plannerContext is null");
-    try (DruidPlanner planner = plannerFactory.createPlannerWithContext(plannerContext)) {
+    try {
       this.plannerResult = planner.plan();
     }
-    // we can't collapse catch clauses since SqlPlanningException has type-sensitive constructors.
-    catch (SqlParseException e) {
-      throw new SqlPlanningException(e);
-    }
     catch (ValidationException e) {
       throw new SqlPlanningException(e);
     }
+    finally {
+      // Done with the planner, close it.
+      planner.close();
+      planner = null;
+    }
   }
 
   /**
@@ -376,20 +391,20 @@ public void after(boolean isDone, Throwable thrown)
     });
   }
 
-
   @VisibleForTesting
-  public ValidationResult runAnalyzeResources(AuthenticationResult authenticationResult)
+  public Set<ResourceAction> runAnalyzeResources(AuthenticationResult authenticationResult)
   {
-    return validate(authenticationResult);
+    validate(authenticationResult);
+    return getRequiredResourceActions();
   }
 
   public Set<ResourceAction> getRequiredResourceActions()
   {
-    return Preconditions.checkNotNull(validationResult, "validationResult").getResourceActions();
+    return resourceActions;
   }
 
   /**
-   * Cancel all native queries associated to this lifecycle.
+   * Cancel all native queries associated with this lifecycle.
    *
    * This method is thread-safe.
    */
@@ -405,7 +420,7 @@ public void cancel()
     final CopyOnWriteArrayList<String> nativeQueryIds = plannerContext.getNativeQueryIds();
 
     for (String nativeQueryId : nativeQueryIds) {
-      log.debug("canceling native query [%s]", nativeQueryId);
+      log.debug("Canceling native query [%s]", nativeQueryId);
       queryScheduler.cancelQuery(nativeQueryId);
     }
   }
@@ -433,13 +448,22 @@ public void finalizeStateAndEmitLogsAndMetrics(
 
       if (state != State.CANCELLED) {
         if (state == State.DONE) {
-          log.warn("Tried to emit logs and metrics twice for query[%s]!", sqlQueryId());
+          log.warn("Tried to emit logs and metrics twice for query [%s]!", sqlQueryId());
         }
 
         state = State.DONE;
       }
     }
 
+    final Set<ResourceAction> actions;
+    if (planner != null) {
+      actions = getRequiredResourceActions();
+      planner.close();
+      planner = null;
+    } else {
+      actions = null;
+    }
+
     final boolean success = e == null;
     final long queryTimeNs = System.nanoTime() - startNs;
 
@@ -449,10 +473,10 @@ public void finalizeStateAndEmitLogsAndMetrics(
         metricBuilder.setDimension("id", plannerContext.getSqlQueryId());
         metricBuilder.setDimension("nativeQueryIds", plannerContext.getNativeQueryIds().toString());
       }
-      if (validationResult != null) {
+      if (actions != null) {
         metricBuilder.setDimension(
             "dataSource",
-            validationResult.getResourceActions()
+            actions
                             .stream()
                             .map(action -> action.getResource().getName())
                             .collect(Collectors.toList())
@@ -527,7 +551,7 @@ private void transition(final State from, final State to)
       }
       if (state != from) {
         throw new ISE(
-            "Cannot transition from[%s] to[%s] because current state[%s] is not [%s].",
+            "Cannot transition from [%s] to [%s] because current state [%s] is not [%s].",
             from,
             to,
             state,

sql/src/main/java/org/apache/druid/sql/avatica/DruidConnection.java

@@ -141,7 +141,6 @@ public DruidJdbcPreparedStatement createPreparedStatement(
         throw DruidMeta.logFailure(new ISE("Too many open statements, limit is [%,d]", maxStatements));
       }
 
-      @SuppressWarnings("GuardedBy")
       final DruidJdbcPreparedStatement jdbcStmt = new DruidJdbcPreparedStatement(
           this,
           statementId,

sql/src/main/java/org/apache/druid/sql/calcite/planner/CalcitePlanner.java

@@ -36,8 +36,8 @@
 import org.apache.calcite.plan.RelTraitDef;
 import org.apache.calcite.plan.RelTraitSet;
 import org.apache.calcite.plan.volcano.VolcanoPlanner;
+import org.apache.calcite.prepare.BaseDruidSqlValidator;
 import org.apache.calcite.prepare.CalciteCatalogReader;
-import org.apache.calcite.prepare.DruidSqlValidator;
 import org.apache.calcite.rel.RelCollationTraitDef;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelRoot;
@@ -237,7 +237,7 @@ public SqlNode validate(SqlNode sqlNode) throws ValidationException
     final SqlConformance conformance = conformance();
     final CalciteCatalogReader catalogReader = createCatalogReader();
     this.validator =
-        new DruidSqlValidator(operatorTable, catalogReader, typeFactory,
+        new BaseDruidSqlValidator(operatorTable, catalogReader, typeFactory,
             conformance);
     this.validator.setIdentifierExpansion(true);
     try {
@@ -324,7 +324,7 @@ public RelRoot rel(SqlNode sql)
     final CalciteCatalogReader catalogReader =
         createCatalogReader().withSchemaPath(schemaPath);
     final SqlValidator validator =
-        new DruidSqlValidator(operatorTable, catalogReader, typeFactory,
+        new BaseDruidSqlValidator(operatorTable, catalogReader, typeFactory,
             conformance);
     validator.setIdentifierExpansion(true);