gianm commented on code in PR #12845:
URL: https://github.com/apache/druid/pull/12845#discussion_r938978573
##########
pom.xml:
##########
@@ -128,7 +128,7 @@
<repoOrgUrl>https://repository.apache.org/snapshots</repoOrgUrl>
<!-- Allow the handful of flaky tests with transient failures to pass.
-->
- <surefire.rerunFailingTestsCount>3</surefire.rerunFailingTestsCount>
+ <!--
<surefire.rerunFailingTestsCount>3</surefire.rerunFailingTestsCount> -->
Review Comment:
Intentional? My sense is we still have a bunch of flaky unit tests, so I
think we still need this rerunner for now.
##########
sql/src/main/java/org/apache/druid/sql/SqlLifecycleManager.java:
##########
@@ -22,40 +22,48 @@
import com.google.common.collect.ImmutableList;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import org.apache.druid.guice.LazySingleton;
-import org.apache.druid.sql.SqlLifecycle.State;
+import org.apache.druid.server.security.ResourceAction;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Set;
/**
- * This class manages only _authorized_ {@link SqlLifecycle}s submitted via
HTTP,
- * such as {@link org.apache.druid.sql.http.SqlResource}. The main use case of
this class is
- * tracking running queries so that the cancel API can identify the lifecycles
to cancel.
+ * This class manages only <i>authorized</i> {@link DirectStatement}s
submitted via
+ * HTTP, such as {@link org.apache.druid.sql.http.SqlResource}. The main use
case of
+ * this class is tracking running queries so that the cancel API can identify
+ * the statements to cancel.
*
- * This class is thread-safe as there are 2 or more threads that can access
lifecycles at the same time
- * for query running or query canceling.
+ * This class is thread-safe as there are 2 or more threads that can access
+ * statements at the same time for query running or query canceling.
*
- * For managing and canceling native queries, see {@link
org.apache.druid.server.QueryScheduler}.
- * As its name indicates, it also performs resource scheduling for native
queries based on query lanes
+ * For managing and canceling native queries, see
+ * {@link org.apache.druid.server.QueryScheduler}. As its name indicates, it
+ * also performs resource scheduling for native queries based on query lanes
* {@link org.apache.druid.server.QueryLaningStrategy}.
*
* @see org.apache.druid.server.QueryScheduler#cancelQuery(String)
*/
@LazySingleton
public class SqlLifecycleManager
{
+ public interface Cancellable
Review Comment:
I suggest we spell this "Cancelable". Rationale:
First, I suggest we follow this style for spelling of forms of "cancel":
https://docs.microsoft.com/en-us/style-guide/a-z-word-list-term-collections/c/cancel-canceled-canceling-cancellation.
(I follow this in my code; and, even though Druid code isn't 100% consistent
in spelling of these words, I'd like to try to make things more consistent.)
Now, the style guide above doesn't have an opinion on "cancelable" vs
"cancellable". So we need to do our own research. Most Druid comments, code,
etc use American English spellings. I checked various American dictionaries and
both seem acceptable. I checked Google's Ngram Viewer and the situation seems…
complicated. "Cancellable" was historically more popular, but has been falling
in popularity relative to "cancelable", which is now (slightly) more popular:
https://books.google.com/ngrams/graph?content=cancelable%2Ccancellable&year_start=1800&year_end=2019&corpus=28&smoothing=3#.
To me, the combination of recent popularity, and the single-L being more
common than double-L across forms of "cancel" generally, suggests that
"cancelable" is the way to go.
##########
sql/src/main/java/org/apache/druid/sql/AbstractStatement.java:
##########
@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql;
+
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.apache.calcite.tools.ValidationException;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.server.security.Access;
+import org.apache.druid.server.security.AuthorizationUtils;
+import org.apache.druid.server.security.ForbiddenException;
+import org.apache.druid.server.security.ResourceAction;
+import org.apache.druid.sql.calcite.planner.DruidPlanner;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
+import org.apache.druid.sql.calcite.planner.PlannerResult;
+
+import java.io.Closeable;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.Function;
+
+/**
+ * Represents a SQL statement either for preparation or execution.
+ * A statement is given by a lifecycle context and the statement
+ * to execute. See derived classes for actions. Closing the statement
+ * emits logs and metrics for the statement.
+ */
+public abstract class AbstractStatement implements Closeable
+{
+ private static final Logger log = new Logger(AbstractStatement.class);
+
+ protected final SqlToolbox sqlToolbox;
+ protected final SqlQueryPlus queryPlus;
+ protected final SqlExecutionReporter reporter;
+ protected PlannerContext plannerContext;
+ protected Set<ResourceAction> queryResource;
+ protected Set<ResourceAction> resourceActions;
+
+ public AbstractStatement(
+ final SqlToolbox sqlToolbox,
+ final SqlQueryPlus sqlRequest,
+ final String remoteAddress
+ )
+ {
+ this.sqlToolbox = sqlToolbox;
+ this.queryPlus = sqlRequest;
+ this.reporter = new SqlExecutionReporter(this, remoteAddress);
+
+ // Context is modified, not copied.
+ contextWithSqlId(sqlRequest.context())
+ .addDefaultParams(sqlToolbox.defaultQueryConfig.getContext());
+ }
+
+ private QueryContext contextWithSqlId(QueryContext queryContext)
Review Comment:
Looks like this can be static. IMO it's good for anything that can be static
to be static, so readers don't need to wonder if object state is accessed or
modified.
##########
sql/src/main/java/org/apache/druid/sql/AbstractStatement.java:
##########
@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql;
+
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.apache.calcite.tools.ValidationException;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.server.security.Access;
+import org.apache.druid.server.security.AuthorizationUtils;
+import org.apache.druid.server.security.ForbiddenException;
+import org.apache.druid.server.security.ResourceAction;
+import org.apache.druid.sql.calcite.planner.DruidPlanner;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
+import org.apache.druid.sql.calcite.planner.PlannerResult;
+
+import java.io.Closeable;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.Function;
+
+/**
+ * Represents a SQL statement either for preparation or execution.
+ * A statement is given by a lifecycle context and the statement
+ * to execute. See derived classes for actions. Closing the statement
+ * emits logs and metrics for the statement.
+ */
+public abstract class AbstractStatement implements Closeable
+{
+ private static final Logger log = new Logger(AbstractStatement.class);
+
+ protected final SqlToolbox sqlToolbox;
+ protected final SqlQueryPlus queryPlus;
+ protected final SqlExecutionReporter reporter;
+ protected PlannerContext plannerContext;
+ protected Set<ResourceAction> queryResource;
+ protected Set<ResourceAction> resourceActions;
+
+ public AbstractStatement(
+ final SqlToolbox sqlToolbox,
+ final SqlQueryPlus sqlRequest,
+ final String remoteAddress
+ )
+ {
+ this.sqlToolbox = sqlToolbox;
+ this.queryPlus = sqlRequest;
+ this.reporter = new SqlExecutionReporter(this, remoteAddress);
+
+ // Context is modified, not copied.
+ contextWithSqlId(sqlRequest.context())
+ .addDefaultParams(sqlToolbox.defaultQueryConfig.getContext());
+ }
+
+ private QueryContext contextWithSqlId(QueryContext queryContext)
+ {
+ // "bySegment" results are never valid to use with SQL because the result
format is incompatible
+ // so, overwrite any user specified context to avoid exceptions down the
line
+
+ if (queryContext.removeUserParam(QueryContexts.BY_SEGMENT_KEY) != null) {
+ log.warn("'bySegment' results are not supported for SQL queries,
ignoring query context parameter");
+ }
+ queryContext.addDefaultParam(PlannerContext.CTX_SQL_QUERY_ID,
UUID.randomUUID().toString());
+ return queryContext;
+ }
+
+ public String sqlQueryId()
+ {
+ return queryPlus.context().getAsString(PlannerContext.CTX_SQL_QUERY_ID);
+ }
+
+ /**
+ * Validate SQL query and authorize against any datasources or views which
+ * will take part in the query. Must be called by the API methods, not
+ * directly.
+ */
+ protected void validate(DruidPlanner planner)
+ {
+ plannerContext = planner.getPlannerContext();
+ plannerContext.setAuthenticationResult(queryPlus.authResult());
+ plannerContext.setParameters(queryPlus.parameters());
+ try {
+ planner.validate();
+ }
+ // 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);
+ }
+ }
+
+ /**
+ * Authorize the query using the authorizer provided, and an option to
authorize
+ * context variables as well as query resources.
+ */
+ protected void authorize(
+ DruidPlanner planner,
+ Function<Set<ResourceAction>, Access> authorizer
+ )
+ {
+ boolean authorizeContextParams =
sqlToolbox.authConfig.authorizeQueryContextParams();
+
+ // Authentication is done by the planner using the function provided
+ // here. The planner ensures that this step is done before planning.
+ Access authorizationResult = planner.authorize(authorizer,
authorizeContextParams);
+ if (!authorizationResult.isAllowed()) {
+ throw new ForbiddenException(authorizationResult.toString());
+ }
+
+ queryResource = planner.resourceActions(false);
Review Comment:
Why do we need to pass an explicit `false` here, instead of
`authorizeContextParams`? Seems sketchy, so there should be a comment
explaining this.
I also suggest having the names be more aligned, since they're very similar
things. So, instead of `resourceActions` and `queryResource`, something like
`resourceActions` and `resourceActionsWithoutContext`.
##########
sql/src/main/java/org/apache/druid/sql/PreparedStatement.java:
##########
@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql;
+
+import org.apache.calcite.avatica.remote.TypedValue;
+import org.apache.calcite.tools.ValidationException;
+import org.apache.druid.sql.calcite.planner.DruidPlanner;
+import org.apache.druid.sql.calcite.planner.PrepareResult;
+
+import java.util.List;
+
+/**
+ * Statement for the JDBC prepare-once, execute many model.
+ */
+public class PreparedStatement extends AbstractStatement
+{
+ private final SqlQueryPlus originalRequest;
+ private PrepareResult prepareResult;
+
+ public PreparedStatement(
+ final SqlToolbox lifecycleToolbox,
+ final SqlQueryPlus queryPlus
+ )
+ {
+ super(lifecycleToolbox, queryPlus, null);
+ this.originalRequest = queryPlus;
+ }
+
+ /**
+ * Prepare the query lifecycle for execution, without completely planning
into
+ * something that is executable, but including some initial parsing and
+ * validation, to support prepared statements via JDBC.
+ * <p>
+ * Note that, per JDBC convention, the prepare step does not provide
+ * parameter values: those are provided later during execution and will
generally
+ * vary from one execution to the next.
+ *
+ * <ul>
+ * <li>Create the planner.</li>
+ * <li>Parse the statement.</li>
+ * <li>JDBC does not provide parameter values at prepare time.
+ * They are provided during execution later, where we'll replan the
+ * query to use the <a href="https://github.com/apache/druid/pull/6974";>
+ * "query optimized"</a> structure.</li>
+ * <li>Validate the query against the Druid catalog.</li>
+ * <li>Authorize access to the resources which the query needs.</li>
+ * <li>Return a {@link PrepareResult} which describes the query.</li>
+ * </ul>
+ */
+ public PrepareResult prepare()
+ {
+ try (DruidPlanner planner = sqlToolbox.plannerFactory.createPlanner(
+ queryPlus.sql(),
+ queryPlus.context())) {
+ validate(planner);
+ authorize(planner, authorizer());
Review Comment:
Nevermind, I see that we're already doing that because we return a new
DirectStatement every time `execute` runs. A comment about this would be
useful, in case someone decides to refactor things later.
##########
sql/src/main/java/org/apache/druid/sql/AbstractStatement.java:
##########
@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql;
+
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.apache.calcite.tools.ValidationException;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.server.security.Access;
+import org.apache.druid.server.security.AuthorizationUtils;
+import org.apache.druid.server.security.ForbiddenException;
+import org.apache.druid.server.security.ResourceAction;
+import org.apache.druid.sql.calcite.planner.DruidPlanner;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
+import org.apache.druid.sql.calcite.planner.PlannerResult;
+
+import java.io.Closeable;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.Function;
+
+/**
+ * Represents a SQL statement either for preparation or execution.
+ * A statement is given by a lifecycle context and the statement
+ * to execute. See derived classes for actions. Closing the statement
+ * emits logs and metrics for the statement.
+ */
+public abstract class AbstractStatement implements Closeable
+{
+ private static final Logger log = new Logger(AbstractStatement.class);
+
+ protected final SqlToolbox sqlToolbox;
+ protected final SqlQueryPlus queryPlus;
+ protected final SqlExecutionReporter reporter;
+ protected PlannerContext plannerContext;
+ protected Set<ResourceAction> queryResource;
+ protected Set<ResourceAction> resourceActions;
+
+ public AbstractStatement(
+ final SqlToolbox sqlToolbox,
+ final SqlQueryPlus sqlRequest,
+ final String remoteAddress
+ )
+ {
+ this.sqlToolbox = sqlToolbox;
+ this.queryPlus = sqlRequest;
Review Comment:
Any reason to use different variable names for the constructor param and the
field?
##########
sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidTypeSystem.java:
##########
@@ -29,7 +29,8 @@ public class DruidTypeSystem implements RelDataTypeSystem
public static final DruidTypeSystem INSTANCE = new DruidTypeSystem();
/**
- * Druid uses millisecond precision for timestamps internally. This is also
the default at the SQL layer.
+ * Druid uses millisecond precision for timestamps internally. This is also
+ * the default at the SQL layer.
Review Comment:
No reason to break this line. Our preferred line width is 120 characters,
which this line fits in. (See comments in `checkstyle.xml`, and the
`org.eclipse.jdt.core.formatter.lineSplit` setting in `eclipse_formatting.xml`.)
##########
processing/src/main/java/org/apache/druid/query/QueryContext.java:
##########
@@ -68,6 +68,19 @@ public QueryContext(@Nullable Map<String, Object> userParams)
invalidateMergedParams();
}
+ public QueryContext(QueryContext from, @Nullable Map<String, Object>
userParams)
+ {
+ this.defaultParams = new TreeMap<>(from.defaultParams);
+ this.userParams = userParams == null ? new TreeMap<>() : new
TreeMap<>(userParams);
+ this.systemParams = new TreeMap<>(from.systemParams);
+ invalidateMergedParams();
+ }
+
+ public QueryContext withOverrides(Map<String, Object> overrides)
Review Comment:
This could use some javadocs too, now that certain methods entirely replace
userParams with the provided map and some methods use the provided map to
override specific keys. People need to be certain of which behavior they're
going to get.
##########
sql/src/main/resources/saffron.properties:
##########
@@ -0,0 +1,28 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#-------------------------------------------------------------
+#
+# Properties for Calcite (formerly known as "Saffron").
Review Comment:
I've always wondered why it was called `saffron.properties`.
##########
sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidPlanner.java:
##########
@@ -244,7 +237,13 @@ public PrepareResult prepare() throws ValidationException
Preconditions.checkState(state == State.VALIDATED);
rootQueryRel = planner.rel(validatedQueryNode);
+ doPrepare(null);
+ state = State.PREPARED;
+ return prepareResult;
+ }
+ private void doPrepare(QueryMaker queryMaker) throws ValidationException
Review Comment:
Since `queryMaker` here is nullable, it's nice to include `@Nullable`.
##########
sql/src/main/java/org/apache/druid/sql/PreparedStatement.java:
##########
@@ -0,0 +1,106 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql;
+
+import org.apache.calcite.avatica.remote.TypedValue;
+import org.apache.calcite.tools.ValidationException;
+import org.apache.druid.sql.calcite.planner.DruidPlanner;
+import org.apache.druid.sql.calcite.planner.PrepareResult;
+
+import java.util.List;
+
+/**
+ * Statement for the JDBC prepare-once, execute many model.
+ */
+public class PreparedStatement extends AbstractStatement
+{
+ private final SqlQueryPlus originalRequest;
+ private PrepareResult prepareResult;
+
+ public PreparedStatement(
+ final SqlToolbox lifecycleToolbox,
+ final SqlQueryPlus queryPlus
+ )
+ {
+ super(lifecycleToolbox, queryPlus, null);
+ this.originalRequest = queryPlus;
+ }
+
+ /**
+ * Prepare the query lifecycle for execution, without completely planning
into
+ * something that is executable, but including some initial parsing and
+ * validation, to support prepared statements via JDBC.
+ * <p>
+ * Note that, per JDBC convention, the prepare step does not provide
+ * parameter values: those are provided later during execution and will
generally
+ * vary from one execution to the next.
+ *
+ * <ul>
+ * <li>Create the planner.</li>
+ * <li>Parse the statement.</li>
+ * <li>JDBC does not provide parameter values at prepare time.
+ * They are provided during execution later, where we'll replan the
+ * query to use the <a href="https://github.com/apache/druid/pull/6974";>
+ * "query optimized"</a> structure.</li>
+ * <li>Validate the query against the Druid catalog.</li>
+ * <li>Authorize access to the resources which the query needs.</li>
+ * <li>Return a {@link PrepareResult} which describes the query.</li>
+ * </ul>
+ */
+ public PrepareResult prepare()
+ {
+ try (DruidPlanner planner = sqlToolbox.plannerFactory.createPlanner(
+ queryPlus.sql(),
+ queryPlus.context())) {
+ validate(planner);
+ authorize(planner, authorizer());
Review Comment:
Shall we re-authorize the statement every time it runs? That way, if
someone's permissions are adjusted such that they can no longer run the query,
they don't have a backdoor way to keep running it.
##########
sql/src/main/java/org/apache/druid/sql/DirectStatement.java:
##########
@@ -0,0 +1,179 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql;
+
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.guava.Sequence;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.QueryInterruptedException;
+import org.apache.druid.sql.SqlLifecycleManager.Cancellable;
+import org.apache.druid.sql.calcite.planner.DruidPlanner;
+import org.apache.druid.sql.calcite.planner.PlannerResult;
+import org.apache.druid.sql.calcite.planner.PrepareResult;
+
+import java.util.concurrent.CopyOnWriteArrayList;
+
+/**
+ * Lifecycle for direct SQL statement execution, which means that the query
+ * is planned and executed in a single step, with no "prepare" step.
+ * Callers need call only:
+ * <ul>
+ * <li>{@link #execute()} to execute the query. The caller must close
+ * the returned {@code Sequence}.</li>
+ * <li>{@link #close()} to report metrics, or {@link #closeQuietly()}
+ * otherwise.</li>
+ * </ul>
+ * <p>
+ * The {@link #cancel()} method may be called from any thread and cancels
+ * the query.
+ * <p>
+ * All other methods are optional and are generally for introspection.
+ * <p>
+ * The class supports two threading models. In the simple case, the same
+ * thread creates this object and executes the query. In the split model,
+ * a request thread creates this object and plans the query. A separate
+ * response thread consumes results and performs any desired logging, etc.
+ * The object is transferred between threads, with no overlapping access.
+ * <p>
+ * As statement holds no resources and need not be called. Only the
+ * {@code Sequence} returned from {@link #execute()} need be closed.
+ * <p>
+ * Use this class for tests and JDBC execution. Use the HTTP variant,
+ * {@link HttpStatement} for HTTP requests.
+ */
+public class DirectStatement extends AbstractStatement implements Cancellable
+{
+ private static final Logger log = new Logger(DirectStatement.class);
+
+ protected PrepareResult prepareResult;
+ protected PlannerResult plannerResult;
+ private volatile boolean cancelled;
Review Comment:
Canceled (spelling (sorry))
##########
processing/src/main/java/org/apache/druid/query/QueryContext.java:
##########
@@ -68,6 +68,19 @@ public QueryContext(@Nullable Map<String, Object> userParams)
invalidateMergedParams();
}
+ public QueryContext(QueryContext from, @Nullable Map<String, Object>
userParams)
Review Comment:
Looks like this drops `from.userParams` and replaces it with `userParams`.
Before I read the code, my guess was that it would combine them in the manner
of `QueryContexts.override`. So, IMO, this constructor needs some javadoc
explaining what it does.
Matter of taste: it would also make sense to me for it to be an instance
method, rather than constructor, so usage would be like `new
QueryContext(otherContext).withUserParams(userParams)` instead of `new
QueryContext(otherContext, userParams)`.
Up to you, really, what you think is best. The only thing I'm asking for
here is that there be some additional clarification about the behavior.
##########
sql/src/main/java/org/apache/druid/sql/AbstractStatement.java:
##########
@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql;
+
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.apache.calcite.tools.ValidationException;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.server.security.Access;
+import org.apache.druid.server.security.AuthorizationUtils;
+import org.apache.druid.server.security.ForbiddenException;
+import org.apache.druid.server.security.ResourceAction;
+import org.apache.druid.sql.calcite.planner.DruidPlanner;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
+import org.apache.druid.sql.calcite.planner.PlannerResult;
+
+import java.io.Closeable;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.Function;
+
+/**
+ * Represents a SQL statement either for preparation or execution.
+ * A statement is given by a lifecycle context and the statement
+ * to execute. See derived classes for actions. Closing the statement
+ * emits logs and metrics for the statement.
+ */
+public abstract class AbstractStatement implements Closeable
+{
+ private static final Logger log = new Logger(AbstractStatement.class);
+
+ protected final SqlToolbox sqlToolbox;
+ protected final SqlQueryPlus queryPlus;
+ protected final SqlExecutionReporter reporter;
+ protected PlannerContext plannerContext;
+ protected Set<ResourceAction> queryResource;
+ protected Set<ResourceAction> resourceActions;
+
+ public AbstractStatement(
+ final SqlToolbox sqlToolbox,
+ final SqlQueryPlus sqlRequest,
+ final String remoteAddress
+ )
+ {
+ this.sqlToolbox = sqlToolbox;
+ this.queryPlus = sqlRequest;
+ this.reporter = new SqlExecutionReporter(this, remoteAddress);
+
+ // Context is modified, not copied.
+ contextWithSqlId(sqlRequest.context())
+ .addDefaultParams(sqlToolbox.defaultQueryConfig.getContext());
+ }
+
+ private QueryContext contextWithSqlId(QueryContext queryContext)
+ {
+ // "bySegment" results are never valid to use with SQL because the result
format is incompatible
+ // so, overwrite any user specified context to avoid exceptions down the
line
+
+ if (queryContext.removeUserParam(QueryContexts.BY_SEGMENT_KEY) != null) {
+ log.warn("'bySegment' results are not supported for SQL queries,
ignoring query context parameter");
+ }
+ queryContext.addDefaultParam(PlannerContext.CTX_SQL_QUERY_ID,
UUID.randomUUID().toString());
+ return queryContext;
+ }
+
+ public String sqlQueryId()
+ {
+ return queryPlus.context().getAsString(PlannerContext.CTX_SQL_QUERY_ID);
+ }
+
+ /**
+ * Validate SQL query and authorize against any datasources or views which
+ * will take part in the query. Must be called by the API methods, not
+ * directly.
+ */
+ protected void validate(DruidPlanner planner)
+ {
+ plannerContext = planner.getPlannerContext();
+ plannerContext.setAuthenticationResult(queryPlus.authResult());
+ plannerContext.setParameters(queryPlus.parameters());
+ try {
+ planner.validate();
+ }
+ // 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);
+ }
+ }
+
+ /**
+ * Authorize the query using the authorizer provided, and an option to
authorize
+ * context variables as well as query resources.
+ */
+ protected void authorize(
+ DruidPlanner planner,
+ Function<Set<ResourceAction>, Access> authorizer
+ )
+ {
+ boolean authorizeContextParams =
sqlToolbox.authConfig.authorizeQueryContextParams();
+
+ // Authentication is done by the planner using the function provided
+ // here. The planner ensures that this step is done before planning.
+ Access authorizationResult = planner.authorize(authorizer,
authorizeContextParams);
+ if (!authorizationResult.isAllowed()) {
+ throw new ForbiddenException(authorizationResult.toString());
+ }
+
+ queryResource = planner.resourceActions(false);
+ resourceActions = planner.resourceActions(authorizeContextParams);
+ }
+
+ /**
+ * Resource authorizer based on the authentication result
+ * provided earlier.
+ */
+ protected Function<Set<ResourceAction>, Access> authorizer()
+ {
+ return resourceActions ->
+ AuthorizationUtils.authorizeAllResourceActions(
+ queryPlus.authResult(),
+ resourceActions,
+ sqlToolbox.plannerFactory.getAuthorizerMapper()
+ );
+ }
+
+ /**
+ * Plan the query, which also produces the sequence that runs
+ * the query.
+ */
+ protected PlannerResult plan(DruidPlanner planner)
+ {
+ try {
+ return planner.plan();
+ }
+ catch (ValidationException e) {
+ throw new SqlPlanningException(e);
+ }
+ }
+
+ /**
+ * Return the datasource and table resources for this
+ * statement.
+ */
+ public Set<ResourceAction> resources()
+ {
+ return queryResource;
+ }
+
+ public Set<ResourceAction> allResources()
+ {
+ return resourceActions;
+ }
+
+ public SqlQueryPlus sqlRequest()
+ {
+ return queryPlus;
+ }
+
+ public SqlExecutionReporter reporter()
+ {
+ return reporter;
+ }
+
+ /**
+ * Releases resources and emits logs and metrics as defined the
+ * associated reporter.
+ */
+ @Override
+ public void close()
+ {
+ closeQuietly();
+ reporter.emit();
Review Comment:
I'm wondering if there should be a try/finally here. Can `closeQuietly`
throw errors? If so, should we still call `reporter.emit()` in that case?
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
