[fix](flight-sql) Support prepared query parameter binding on branch-4.1 (#68768)

### What problem does this PR solve?

Flight SQL prepared queries containing `?` fail because Prepare rejects
placeholders and `acceptPutPreparedStatementQuery` returns
`UNIMPLEMENTED`. For example, ADBC cannot bind an integer to `SELECT
CAST(? AS BIGINT)` or execute a range predicate with two bound values.

This change implements single-row typed parameter binding on branch-4.1.
Prepare analyzes explicit casts and comparison columns to advertise
parameter types and a best-effort result schema. DoPut converts Arrow
values into detached Nereids literals, and execution binds those
literals to a fresh statement context. Result schemas are reanalyzed
after binding. Values are never interpolated into SQL text.

Arrow JDBC receives typed parameter metadata and a result schema before
binding, so it can construct parameter vectors and classify SELECT
correctly. Types that cannot be inferred remain unknown; JDBC callers
should use an explicit CAST for queries such as `SELECT ?`.
Timezone-bearing timestamp uploads preserve their UTC instants as
TIMESTAMPTZ literals; timestamps without a timezone retain wall-clock
semantics. Unsupported parameter vector types are not advertised as
bindable. Invalid SQL returns `INVALID_ARGUMENT` consistently through
Prepare and GetSchema.

Failed rebinds invalidate previous values. Handle ownership, concurrent
closure, upload versions, and per-binding/session memory limits protect
retained state. Multi-row uploads and unsupported types fail explicitly.
Queries requiring forwarding to the master FE reject bound parameters
because the forwarding protocol cannot carry typed bindings.

The PR also removes an obsolete duplicate JWT version property,
preserving the effective dependency version. The target branch already
contains the dependency declaration cleanup.

### Release note

Support single-row scalar parameter binding for Flight SQL prepared
queries, including repeated execution and NULL values.

### Check List (For Author)

- Test
    - [x] Regression test
    - [x] Unit Test
    - [x] Manual test
- Behavior changed:
- [x] Yes. Parameterized Flight SQL queries can bind and execute
supported scalar values.
- Does this need documentation?
    - [x] No. This implements the existing prepared-query protocol.

Validation:
- Full FE Checkstyle passed (`mvn -B clean checkstyle:check`).
- Focused compilation of all 11 changed Java source/test files passed
with `javac --release 8` using existing FE dependency artifacts. Java
8-compatible Guava copies preserve immutable parameter snapshots.
- 23 new tests and 33 existing Flight SQL schema tests passed, including
invalid-SQL status checks over Flight RPC and connection recovery.
- The new regression suite passed against a local FE service test
harness and a real BE, including actual JDBC `PreparedStatement`
execution, result fetching, rebinding, NULL, mixed types, range
predicates, and recovery after invalid uploads. TIMESTAMPTZ predicates
are checked under UTC, Asia/Shanghai and America/New_York, including a
daylight-saving transition.
- ADBC 1.11.0 with PyArrow 23.0.1 passed 18 concurrent connections and
540 parameterized queries, with result verification plus NULL and
connection-reuse checks. A further 72 timestamp result checks cover all
four Arrow units, three session timezones, two input timezone
annotations, negative epochs and NULL values.

The local integration harness uses one FE service and one real BE; a
full multi-FE cluster run and full FE build are left to CI.

### Check List (For Reviewer who merge this PR)

- [ ] Confirm the release note
- [ ] Confirm test cases
- [ ] Confirm document
- [ ] Add branch pick label
This commit is contained in:
Gabriel authored and GitHub committed 2026-10-09 14:37:12 +08:00
1 parent da03ee7824
commit bd4ce9fbc0
12 files changed
+1655 -65

No files matched your search

@@ -813,7 +813,10 @@ public class ExpressionAnalyzer extends SubExprAnalyzer<ExpressionRewriteContext
private void registerPlaceholderIdToSlot(ComparisonPredicate cp,
ExpressionRewriteContext context, Expression left, Expression right) {
if (ConnectContext.get() != null
&& ConnectContext.get().getCommand() == MysqlCommand.COM_STMT_EXECUTE) {
&& (ConnectContext.get().getCommand() == MysqlCommand.COM_STMT_EXECUTE
// Flight Prepare needs the bound column type before comparison coercion erases that constraint.
|| (ConnectContext.get().getConnectType() == ConnectContext.ConnectType.ARROW_FLIGHT_SQL
&& context.cascadesContext.getStatementContext().isPrepareStage()))) {
// Used to replace expression in ShortCircuit plan
if (cp.right() instanceof Placeholder && left instanceof SlotReference) {
PlaceholderId id = ((Placeholder) cp.right()).getPlaceholderId();
@@ -87,6 +87,7 @@ import org.apache.doris.transaction.TransactionEntry;
import org.apache.doris.transaction.TransactionStatus;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import lombok.Getter;
@@ -156,12 +157,19 @@ public class ConnectContext {
// for arrow flight
protected volatile String peerIdentity;
private final Map<String, PreparedQuery> preparedQuerys = new HashMap<>();
private static final long MAX_FLIGHT_PARAMETER_BYTES = 16L * 1024 * 1024;
private long flightParameterBytes;
private boolean flightPreparedQueriesClosed;
private static class PreparedQuery {
private final String sql;
private final String catalog;
private final String database;
private final Schema schema;
private String catalog;
private String database;
private Schema schema;
private int parameterCount;
private List<Literal> parameters;
private long bindingVersion;
private long parameterBytes;
private PreparedQuery(String sql, String catalog, String database, Schema schema) {
this.sql = sql;
@@ -427,7 +435,7 @@ public class ConnectContext {
}
resetSessionVariable();
userVars = new HashMap<>();
preparedQuerys.clear();
clearPreparedQueries();
preparedStatementContextMap.clear();
runningQuery = null;
queryId = null;
@@ -924,35 +932,157 @@ public class ConnectContext {
this.loginTime = System.currentTimeMillis();
}
public synchronized void addPreparedQuery(String preparedStatementId, String preparedQuery) {
addPreparedQuery(preparedStatementId, preparedQuery, null);
}
public synchronized void addPreparedQuery(String preparedStatementId, String preparedQuery, Schema schema) {
preparedQuerys.put(preparedStatementId,
new PreparedQuery(preparedQuery, getDefaultCatalog(), getDatabase(), schema));
}
public synchronized Schema getPreparedQuerySchema(String preparedStatementId) {
PreparedQuery query = preparedQuerys.get(preparedStatementId);
return query == null ? null : query.schema;
}
public synchronized String getPreparedQuery(String preparedStatementId) {
PreparedQuery query = preparedQuerys.get(preparedStatementId);
if (query == null) {
return null;
public void addPreparedQuery(String preparedStatementId, String preparedQuery) {
synchronized (preparedQuerys) {
addPreparedQuery(preparedStatementId, preparedQuery, null);
}
// A handle must not execute unqualified SQL in a different namespace than its advertised schema.
if (!Objects.equals(query.catalog, getDefaultCatalog()) || !Objects.equals(query.database, getDatabase())) {
preparedQuerys.remove(preparedStatementId);
return null;
}
return query.sql;
}
public synchronized void removePreparedQuery(String preparedStatementId) {
preparedQuerys.remove(preparedStatementId);
public void addPreparedQuery(String preparedStatementId, String preparedQuery, Schema schema) {
synchronized (preparedQuerys) {
if (flightPreparedQueriesClosed) {
throw new IllegalStateException("Flight SQL session is closed");
}
removePreparedQuery(preparedStatementId);
preparedQuerys.put(preparedStatementId,
new PreparedQuery(preparedQuery, getDefaultCatalog(), getDatabase(), schema));
}
}
public void addPreparedQuery(String id, String sql, Schema schema, int parameterCount) {
synchronized (preparedQuerys) {
addPreparedQuery(id, sql, schema);
preparedQuerys.get(id).parameterCount = parameterCount;
}
}
public int getPreparedQueryParameterCount(String id) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.get(id);
return query == null ? -1 : query.parameterCount;
}
}
public List<Literal> getPreparedQueryParameters(String id) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.get(id);
return query == null ? null : query.parameters;
}
}
public long beginPreparedQueryBinding(String id) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.get(id);
if (query == null) {
return -1;
}
flightParameterBytes -= query.parameterBytes;
query.parameterBytes = 0;
query.parameters = null;
return ++query.bindingVersion;
}
}
public boolean isPreparedQueryBindingCurrent(String id, long version) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.get(id);
return query != null && query.bindingVersion == version;
}
}
public void setPreparedQueryParameters(String id, List<Literal> parameters, Schema schema) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.get(id);
if (query == null) {
throw new IllegalStateException("Prepared statement expired");
}
long bytes = parameters == null ? 0 : parameters.size() * 128L;
if (parameters != null) {
for (Literal parameter : parameters) {
if (parameter.getValue() instanceof String) {
bytes += 2L * ((String) parameter.getValue()).length();
}
}
}
// A per-upload limit alone allows many handles to retain unbounded parameter memory in one session.
if (flightParameterBytes - query.parameterBytes + bytes > MAX_FLIGHT_PARAMETER_BYTES) {
throw new IllegalArgumentException("Prepared query parameters exceed the 16 MiB session limit");
}
flightParameterBytes += bytes - query.parameterBytes;
query.parameterBytes = bytes;
// Keep a detached immutable snapshot using a Java 8-compatible copy implementation.
query.parameters = parameters == null ? null : ImmutableList.copyOf(parameters);
query.schema = schema;
}
}
public boolean setPreparedQueryParameters(String id, List<Literal> parameters, Schema schema, long version) {
synchronized (preparedQuerys) {
if (!isPreparedQueryBindingCurrent(id, version)) {
return false;
}
setPreparedQueryParameters(id, parameters, schema);
return true;
}
}
public boolean refreshPreparedQueryNamespace(String id) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.get(id);
if (query == null) {
return false;
}
if (query.parameterCount == 0) {
query.catalog = getDefaultCatalog();
query.database = getDatabase();
}
return true;
}
}
public Schema getPreparedQuerySchema(String preparedStatementId) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.get(preparedStatementId);
return query == null ? null : query.schema;
}
}
public String getPreparedQuery(String preparedStatementId) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.get(preparedStatementId);
if (query == null) {
return null;
}
// A handle must not execute unqualified SQL in a different namespace than its advertised schema.
if (!Objects.equals(query.catalog, getDefaultCatalog()) || !Objects.equals(query.database, getDatabase())) {
removePreparedQuery(preparedStatementId);
return null;
}
return query.sql;
}
}
public void closePreparedQueries() {
synchronized (preparedQuerys) {
flightPreparedQueriesClosed = true;
clearPreparedQueries();
}
}
public void clearPreparedQueries() {
synchronized (preparedQuerys) {
preparedQuerys.clear();
flightParameterBytes = 0;
}
}
public void removePreparedQuery(String preparedStatementId) {
synchronized (preparedQuerys) {
PreparedQuery query = preparedQuerys.remove(preparedStatementId);
if (query != null) {
flightParameterBytes -= query.parameterBytes;
}
}
}
public void setRunningQuery(String runningQuery) {
@@ -1183,6 +1183,12 @@ public class StmtExecutor {
}
private void forwardToMaster() throws Exception {
// COM_QUERY forwarding carries SQL text but not Flight's typed placeholder bindings.
if (context.getConnectType() == ConnectType.ARROW_FLIGHT_SQL
&& !statementContext.getIdToPlaceholderRealExpr().isEmpty()) {
throw new UserException("Flight SQL queries with bound parameters cannot be forwarded; "
+ "connect to master FE");
}
masterOpExecutor = new MasterOpExecutor(originStmt, context, redirectStatus, isQuery());
if (LOG.isDebugEnabled()) {
LOG.debug("need to transfer to Master. stmt: {}", context.getStmtId());
@@ -27,6 +27,7 @@ import org.apache.doris.common.util.DebugUtil;
import org.apache.doris.common.util.Util;
import org.apache.doris.mysql.MysqlCommand;
import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.QueryState.MysqlStateType;
import org.apache.doris.qe.StmtExecutor;
@@ -186,10 +187,12 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
final StreamListener<Result> listener) {
executorService.submit(() -> {
try {
String[] handleParts = request.getPreparedStatementHandle().toStringUtf8().split(":");
String executedPeerIdentity = handleParts[0];
String preparedStatementId = handleParts[1];
flightSessionsManager.getConnectContext(executedPeerIdentity).removePreparedQuery(preparedStatementId);
String preparedStatementId = preparedStatementId(context, request.getPreparedStatementHandle());
ConnectContext connection = flightSessionsManager.getConnectContext(context.peerIdentity());
// Closing one handle serializes with its execution; closing a session can still cancel immediately.
synchronized (connection) {
connection.removePreparedQuery(preparedStatementId);
}
} catch (final Throwable e) {
listener.onError(e);
return;
@@ -211,6 +214,11 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
private Pair<FlightInfo, StmtExecutor> executeQueryStatementLocked(String peerIdentity,
ConnectContext connectContext, String query,
final FlightDescriptor descriptor) {
return executeQueryStatementLocked(peerIdentity, connectContext, query, descriptor, null);
}
private Pair<FlightInfo, StmtExecutor> executeQueryStatementLocked(String peerIdentity,
ConnectContext connectContext, String query, FlightDescriptor descriptor, List<Literal> parameters) {
try {
Preconditions.checkState(null != connectContext);
Preconditions.checkState(!query.isEmpty());
@@ -222,7 +230,11 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
connectContext.getFlightSqlChannel().reset();
connectContext.clearFlightSqlEndpointsLocations();
try (FlightSqlConnectProcessor flightSQLConnectProcessor = new FlightSqlConnectProcessor(connectContext)) {
flightSQLConnectProcessor.handleQuery(query);
if (parameters == null) {
flightSQLConnectProcessor.handleQuery(query);
} else {
flightSQLConnectProcessor.handleQuery(query, parameters);
}
if (connectContext.getState().getStateType() == MysqlStateType.ERR) {
throw new RuntimeException("after executeQueryStatement handleQuery");
}
@@ -348,8 +360,10 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
ConnectContext connection = flightSessionsManager.getConnectContext(context.peerIdentity());
synchronized (connection) {
Pair<String, Schema> prepared = preparedQuery(connection, context, command);
String preparedId = preparedStatementId(context, command.getPreparedStatementHandle());
Pair<FlightInfo, StmtExecutor> result = executeQueryStatementLocked(
context.peerIdentity(), connection, prepared.getLeft(), descriptor);
context.peerIdentity(), connection, prepared.getLeft(), descriptor,
connection.getPreparedQueryParameters(preparedId));
FlightInfo info = result.getLeft();
String id = command.getPreparedStatementHandle().toStringUtf8()
.substring(context.peerIdentity().length() + 1);
@@ -361,7 +375,8 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
? FlightSqlQuerySchema.matchesExecutionSchema(prepared.getRight(), info.getSchema(),
((LogicalPlanAdapter) executor.getParsedStmt()).getColLabels())
: prepared.getRight().equals(info.getSchema());
if (!matches) {
// Session teardown can invalidate the handle during execution without taking the connection monitor.
if (!matches || !connection.refreshPreparedQueryNamespace(id)) {
connection.removePreparedQuery(id);
try {
if (executor != null) {
@@ -377,9 +392,6 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
throw CallStatus.NOT_FOUND.withDescription("Prepared statement schema changed; prepare again")
.toRuntimeException();
}
// A successful USE/SWITCH in this statement may change its own namespace. Rebind
// only this handle; unrelated handles must still reject an external namespace change.
connection.addPreparedQuery(id, prepared.getLeft(), prepared.getRight());
return info;
}
}
@@ -396,25 +408,27 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
final CallContext context, final FlightDescriptor descriptor) {
ConnectContext connection = flightSessionsManager.getConnectContext(context.peerIdentity());
synchronized (connection) {
String id = preparedStatementId(context, command.getPreparedStatementHandle());
String query = connection.getPreparedQuery(id);
if (query != null && connection.getPreparedQueryParameterCount(id) > 0
&& connection.getPreparedQueryParameters(id) == null) {
// Metadata discovery does not execute a query and must work before values are bound.
return new SchemaResult(prepareQuerySchema(connection, query).first);
}
return new SchemaResult(preparedQuery(connection, context, command).getRight());
}
}
private Pair<String, Schema> preparedQuery(ConnectContext connection, CallContext context,
CommandPreparedStatementQuery command) {
String prefix = context.peerIdentity() + ":";
String handle = command.getPreparedStatementHandle().toStringUtf8();
if (!handle.startsWith(prefix)) {
throw CallStatus.INVALID_ARGUMENT.withDescription("Invalid prepared statement handle").toRuntimeException();
}
String id = handle.substring(prefix.length());
String id = preparedStatementId(context, command.getPreparedStatementHandle());
String query = connection.getPreparedQuery(id);
if (query == null) {
throw CallStatus.NOT_FOUND
.withDescription("Prepared statement expired; prepare again in the current namespace")
.toRuntimeException();
}
Schema schema = analyzeQuerySchema(connection, query);
Schema schema = analyzeQuerySchema(connection, query, connection.getPreparedQueryParameters(id));
// Execution reparses SQL using the current session. Never silently replace the schema
// advertised by Prepare when settings such as sql_mode or time_zone change its result.
if (!schema.equals(connection.getPreparedQuerySchema(id))) {
@@ -425,9 +439,22 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
return Pair.of(query, schema);
}
private String preparedStatementId(CallContext context, ByteString handleBytes) {
String prefix = context.peerIdentity() + ":";
String handle = handleBytes.toStringUtf8();
if (!handle.startsWith(prefix) || handle.length() == prefix.length()) {
throw CallStatus.INVALID_ARGUMENT.withDescription("Invalid prepared statement handle").toRuntimeException();
}
return handle.substring(prefix.length());
}
private Schema analyzeQuerySchema(ConnectContext context, String query) {
return analyzeQuerySchema(context, query, null);
}
private Schema analyzeQuerySchema(ConnectContext context, String query, List<Literal> parameters) {
try {
return FlightSqlQuerySchema.analyze(context, query);
return FlightSqlQuerySchema.analyze(context, query, parameters);
} catch (FlightRuntimeException e) {
throw e;
} catch (Exception e) {
@@ -436,6 +463,18 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
}
}
private org.apache.doris.common.Pair<Schema, Schema> prepareQuerySchema(ConnectContext context, String query) {
try {
return FlightSqlQuerySchema.prepare(context, query);
} catch (FlightRuntimeException e) {
throw e;
} catch (Exception e) {
// Invalid SQL must report the same client error through Prepare and GetSchema.
throw CallStatus.INVALID_ARGUMENT.withDescription("Cannot determine query schema: " + e.getMessage())
.withCause(e).toRuntimeException();
}
}
@Override
public void close() throws Exception {
AutoCloseables.close(rootAllocator);
@@ -448,16 +487,18 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
private ActionCreatePreparedStatementResult buildCreatePreparedStatementResult(ByteString handle,
Schema parameterSchema, Schema metaData) {
Preconditions.checkState(!Objects.isNull(metaData));
final ByteString bytes = Objects.isNull(parameterSchema) ? ByteString.EMPTY
: ByteString.copyFrom(serializeMetadata(parameterSchema));
return ActionCreatePreparedStatementResult.newBuilder()
.setDatasetSchema(ByteString.copyFrom(serializeMetadata(metaData))).setParameterSchema(bytes)
.setDatasetSchema(metaData == null ? ByteString.EMPTY
: ByteString.copyFrom(serializeMetadata(metaData)))
.setParameterSchema(bytes)
.setPreparedStatementHandle(handle).build();
}
@Override
public void createPreparedStatement(final ActionCreatePreparedStatementRequest request, final CallContext context,
public void createPreparedStatement(final ActionCreatePreparedStatementRequest request,
final CallContext context,
final StreamListener<Result> listener) {
executorService.submit(() -> {
ConnectContext connectContext = null;
@@ -468,12 +509,15 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
String query = request.getQuery();
// ADBC ExecuteSchema reads this dataset schema directly without calling GetSchema.
// Analyze before registering a handle so failed preparation does not retain a query.
Schema schema = analyzeQuerySchema(connectContext, query);
org.apache.doris.common.Pair<Schema, Schema> prepared = prepareQuerySchema(
connectContext, query);
Schema schema = prepared.first;
preparedStatementId = UUID.randomUUID().toString();
ByteString handle = ByteString.copyFromUtf8(context.peerIdentity() + ":" + preparedStatementId);
Result result = new Result(Any.pack(buildCreatePreparedStatementResult(handle,
new Schema(Collections.emptyList()), schema)).toByteArray());
connectContext.addPreparedQuery(preparedStatementId, query, schema);
prepared.second, schema)).toByteArray());
connectContext.addPreparedQuery(preparedStatementId, query, schema,
prepared.second.getFields().size());
listener.onNext(result);
listener.onCompleted();
}
@@ -530,8 +574,44 @@ public class DorisFlightSqlProducer implements FlightSqlProducer, AutoCloseable
@Override
public Runnable acceptPutPreparedStatementQuery(CommandPreparedStatementQuery command, CallContext context,
FlightStream flightStream, StreamListener<PutResult> ackStream) {
throw CallStatus.UNIMPLEMENTED.withDescription("acceptPutPreparedStatementQuery unimplemented")
.toRuntimeException();
return () -> {
try {
ConnectContext connection = flightSessionsManager.getConnectContext(context.peerIdentity());
String id = preparedStatementId(context, command.getPreparedStatementHandle());
int count;
long version;
synchronized (connection) {
if (connection.getPreparedQuery(id) == null) {
throw CallStatus.NOT_FOUND.withDescription("Prepared statement expired").toRuntimeException();
}
count = connection.getPreparedQueryParameterCount(id);
// A failed rebind must not leave the previous values executable.
version = connection.beginPreparedQueryBinding(id);
if (version < 0) {
throw CallStatus.NOT_FOUND.withDescription("Prepared statement expired").toRuntimeException();
}
}
List<Literal> parameters = FlightSqlParameters.read(flightStream, count);
synchronized (connection) {
String query = connection.getPreparedQuery(id);
if (query == null || !connection.isPreparedQueryBindingCurrent(id, version)) {
throw CallStatus.CANCELLED.withDescription("Prepared statement binding was superseded")
.toRuntimeException();
}
Schema schema = analyzeQuerySchema(connection, query, parameters);
if (!connection.setPreparedQueryParameters(id, parameters, schema, version)) {
throw CallStatus.CANCELLED.withDescription("Prepared statement binding was superseded")
.toRuntimeException();
}
}
ackStream.onCompleted();
} catch (FlightRuntimeException e) {
ackStream.onError(e);
} catch (Exception e) {
ackStream.onError(CallStatus.INVALID_ARGUMENT.withDescription("Cannot bind query parameters: "
+ e.getMessage()).withCause(e).toRuntimeException());
}
};
}
@Override
@@ -18,16 +18,20 @@
package org.apache.doris.service.arrowflight;
import org.apache.doris.analysis.Expr;
import org.apache.doris.analysis.StatementBase;
import org.apache.doris.common.ConnectionException;
import org.apache.doris.common.ErrorCode;
import org.apache.doris.common.ErrorReport;
import org.apache.doris.common.Status;
import org.apache.doris.common.util.DebugUtil;
import org.apache.doris.mysql.MysqlCommand;
import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.proto.InternalService;
import org.apache.doris.proto.Types;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.ConnectProcessor;
import org.apache.doris.qe.SessionVariable;
import org.apache.doris.qe.StmtExecutor;
import org.apache.doris.rpc.BackendServiceProxy;
import org.apache.doris.rpc.RpcException;
@@ -61,6 +65,7 @@ import java.util.concurrent.TimeoutException;
public class FlightSqlConnectProcessor extends ConnectProcessor implements AutoCloseable {
private static final Logger LOG = LogManager.getLogger(FlightSqlConnectProcessor.class);
private Schema arrowSchema;
private List<Literal> parameters;
public FlightSqlConnectProcessor(ConnectContext context) {
super(context);
@@ -100,6 +105,39 @@ public class FlightSqlConnectProcessor extends ConnectProcessor implements AutoC
super.handleQuery(query);
}
public void handleQuery(String query, List<Literal> parameters) throws ConnectionException {
this.parameters = parameters;
try {
handleQuery(query);
} finally {
this.parameters = null;
}
}
@Override
protected List<StatementBase> parseWithFallback(String originStmt, String convertedStmt,
SessionVariable sessionVariable) throws ConnectionException {
List<StatementBase> statements = super.parseWithFallback(originStmt, convertedStmt, sessionVariable);
if (parameters != null && !parameters.isEmpty() && statements != null) {
try {
if (statements.size() != 1 || !(statements.get(0) instanceof LogicalPlanAdapter)) {
throw new IllegalArgumentException("Parameters require a single query statement");
}
// Reparse on every execution so bindings cannot retain planner state from a previous query.
FlightSqlParameters.bind(((LogicalPlanAdapter) statements.get(0)).getStatementContext(), parameters);
} catch (RuntimeException e) {
handleQueryException(e, originStmt, null, null);
for (StatementBase statement : statements) {
if (statement instanceof LogicalPlanAdapter) {
((LogicalPlanAdapter) statement).getStatementContext().close();
}
}
return null;
}
}
return statements;
}
// TODO
// private void handleInitDb() {
// handleInitDb(fullDbName);
@@ -0,0 +1,336 @@
// 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.doris.service.arrowflight;
import org.apache.doris.nereids.StatementContext;
import org.apache.doris.nereids.trees.expressions.Cast;
import org.apache.doris.nereids.trees.expressions.Placeholder;
import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral;
import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral;
import org.apache.doris.nereids.trees.expressions.literal.DateTimeV2Literal;
import org.apache.doris.nereids.trees.expressions.literal.DateV2Literal;
import org.apache.doris.nereids.trees.expressions.literal.DecimalV3Literal;
import org.apache.doris.nereids.trees.expressions.literal.DoubleLiteral;
import org.apache.doris.nereids.trees.expressions.literal.FloatLiteral;
import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
import org.apache.doris.nereids.trees.expressions.literal.SmallIntLiteral;
import org.apache.doris.nereids.trees.expressions.literal.StringLiteral;
import org.apache.doris.nereids.trees.expressions.literal.TimestampTzLiteral;
import org.apache.doris.nereids.trees.expressions.literal.TinyIntLiteral;
import org.apache.doris.nereids.trees.plans.PlaceholderId;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.types.BigIntType;
import org.apache.doris.nereids.types.BooleanType;
import org.apache.doris.nereids.types.DataType;
import org.apache.doris.nereids.types.DateTimeV2Type;
import org.apache.doris.nereids.types.DateV2Type;
import org.apache.doris.nereids.types.DecimalV3Type;
import org.apache.doris.nereids.types.DoubleType;
import org.apache.doris.nereids.types.FloatType;
import org.apache.doris.nereids.types.IntegerType;
import org.apache.doris.nereids.types.NullType;
import org.apache.doris.nereids.types.SmallIntType;
import org.apache.doris.nereids.types.StringType;
import org.apache.doris.nereids.types.TimeStampTzType;
import org.apache.doris.nereids.types.TinyIntType;
import org.apache.arrow.flight.CallStatus;
import org.apache.arrow.flight.FlightRuntimeException;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.vector.DateDayVector;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.TimeStampVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.nio.charset.CharacterCodingException;
import java.nio.charset.StandardCharsets;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/** Converts one parameter row into detached values consumed by Nereids placeholder analysis. */
final class FlightSqlParameters {
private static final long MAX_PARAMETER_BYTES = 1024 * 1024;
static final int MAX_PARAMETERS = 1024;
private FlightSqlParameters() {
}
static Map<PlaceholderId, DataType> inferCastTypes(Plan plan) {
Map<PlaceholderId, DataType> types = new HashMap<>();
plan.foreach(node -> {
((Plan) node).getExpressions().forEach(expression -> expression.foreach(child -> {
if (child instanceof Cast && ((Cast) child).child() instanceof Placeholder) {
// Only the innermost explicit cast constrains the value supplied for this placeholder.
types.put(((Placeholder) ((Cast) child).child()).getPlaceholderId(), ((Cast) child).getDataType());
} else if (child instanceof SubqueryExpr) {
types.putAll(inferCastTypes(((SubqueryExpr) child).getQueryPlan()));
}
}));
});
return types;
}
static boolean supportsParameterType(ArrowType type) {
switch (Types.getMinorTypeForArrowType(type)) {
case NULL:
case BIT:
case TINYINT:
case SMALLINT:
case INT:
case BIGINT:
case FLOAT4:
case FLOAT8:
case VARCHAR:
case DATEDAY:
case TIMESTAMPSEC:
case TIMESTAMPMILLI:
case TIMESTAMPMICRO:
case TIMESTAMPNANO:
case TIMESTAMPSECTZ:
case TIMESTAMPMILLITZ:
case TIMESTAMPMICROTZ:
case TIMESTAMPNANOTZ:
case DECIMAL:
return true;
default:
return false;
}
}
static void bind(StatementContext context, List<Literal> parameters) {
if (parameters == null || parameters.size() != context.getPlaceholders().size()) {
throw CallStatus.INVALID_ARGUMENT.withDescription("Bind all query parameters before execution")
.toRuntimeException();
}
for (int i = 0; i < parameters.size(); i++) {
context.getIdToPlaceholderRealExpr().put(
context.getPlaceholders().get(i).getPlaceholderId(), parameters.get(i));
}
}
static List<Literal> read(FlightStream stream, int parameterCount) {
List<Literal> parameters = null;
FlightRuntimeException failure = null;
while (stream.next()) {
// Consume the upload even after validation fails so clients can receive the final gRPC status.
if (failure != null) {
continue;
}
try {
VectorSchemaRoot root = stream.getRoot();
if (root.getFieldVectors().size() != parameterCount) {
throw invalid("Parameter count does not match the prepared query");
}
if (root.getRowCount() == 0) {
continue;
}
if (root.getRowCount() != 1 || parameters != null) {
throw CallStatus.UNIMPLEMENTED.withDescription("Only one parameter row per binding is supported")
.toRuntimeException();
}
parameters = convert(root);
} catch (FlightRuntimeException e) {
failure = e;
} catch (RuntimeException e) {
failure = invalid("Invalid parameter value: " + e.getMessage());
}
}
if (failure != null) {
throw failure;
}
if (parameters == null) {
if (parameterCount == 0) {
return Collections.emptyList();
}
throw invalid("Parameter upload contains no row");
}
return parameters;
}
static List<Literal> convert(VectorSchemaRoot root) {
if (root.getFieldVectors().size() > MAX_PARAMETERS) {
throw invalid("Too many query parameters (maximum 1024)");
}
long bytes = 0;
List<Literal> parameters = new ArrayList<>();
for (FieldVector vector : root.getFieldVectors()) {
bytes += vector.getBufferSize();
if (bytes > MAX_PARAMETER_BYTES) {
throw invalid("Query parameters exceed the 1 MiB binding limit");
}
parameters.add(literal(vector));
}
return parameters;
}
private static Literal literal(FieldVector vector) {
if (vector.getField().getDictionary() != null) {
throw CallStatus.UNIMPLEMENTED.withDescription("Dictionary encoded parameters are not supported")
.toRuntimeException();
}
DataType type;
switch (vector.getMinorType()) {
case NULL:
type = NullType.INSTANCE;
break;
case BIT:
type = BooleanType.INSTANCE;
break;
case TINYINT:
type = TinyIntType.INSTANCE;
break;
case SMALLINT:
type = SmallIntType.INSTANCE;
break;
case INT:
type = IntegerType.INSTANCE;
break;
case BIGINT:
type = BigIntType.INSTANCE;
break;
case FLOAT4:
type = FloatType.INSTANCE;
break;
case FLOAT8:
type = DoubleType.INSTANCE;
break;
case VARCHAR:
type = StringType.INSTANCE;
break;
case DATEDAY:
type = DateV2Type.INSTANCE;
break;
case TIMESTAMPSEC:
case TIMESTAMPMILLI:
case TIMESTAMPMICRO:
case TIMESTAMPNANO:
type = DateTimeV2Type.of(6);
break;
case TIMESTAMPSECTZ:
case TIMESTAMPMILLITZ:
case TIMESTAMPMICROTZ:
case TIMESTAMPNANOTZ:
// Arrow Java uses TZ vectors even for an empty annotation, which still means wall-clock time.
type = ((ArrowType.Timestamp) vector.getField().getType()).getTimezone().isEmpty()
? DateTimeV2Type.of(6) : TimeStampTzType.of(6);
break;
case DECIMAL:
ArrowType.Decimal decimal = (ArrowType.Decimal) vector.getField().getType();
if (decimal.getScale() < 0 || decimal.getScale() > decimal.getPrecision()) {
throw invalid("Unsupported decimal parameter scale");
}
type = DecimalV3Type.createDecimalV3Type(decimal.getPrecision(), decimal.getScale());
break;
default:
throw CallStatus.UNIMPLEMENTED.withDescription(
"Unsupported query parameter type: " + vector.getField().getType()).toRuntimeException();
}
if (vector.isNull(0)) {
return new NullLiteral(type);
}
Object value = vector.getObject(0);
switch (vector.getMinorType()) {
case BIT: return BooleanLiteral.of((Boolean) value);
case TINYINT: return new TinyIntLiteral(((Number) value).byteValue());
case SMALLINT: return new SmallIntLiteral(((Number) value).shortValue());
case INT: return new IntegerLiteral(((Number) value).intValue());
case BIGINT: return new BigIntLiteral(((Number) value).longValue());
case FLOAT4:
case FLOAT8:
double number = ((Number) value).doubleValue();
if (!Double.isFinite(number)) {
throw invalid("Non-finite floating point parameters are not supported");
}
return type instanceof FloatType ? new FloatLiteral((float) number) : new DoubleLiteral(number);
case VARCHAR:
try {
// Reject malformed UTF-8 instead of silently replacing bytes in a bound predicate.
return new StringLiteral(StandardCharsets.UTF_8.newDecoder()
.decode(ByteBuffer.wrap(((VarCharVector) vector).get(0))).toString());
} catch (CharacterCodingException e) {
throw invalid("String parameter is not valid UTF-8");
}
case DECIMAL: return new DecimalV3Literal((DecimalV3Type) type, (BigDecimal) value);
case DATEDAY:
LocalDate date = LocalDate.ofEpochDay(((DateDayVector) vector).get(0));
checkYear(date.getYear());
return new DateV2Literal(date.getYear(), date.getMonthValue(), date.getDayOfMonth());
default:
long timestamp = ((TimeStampVector) vector).get(0);
ArrowType.Timestamp timestampType = (ArrowType.Timestamp) vector.getField().getType();
long units;
switch (timestampType.getUnit()) {
case SECOND:
units = 1;
break;
case MILLISECOND:
units = 1000;
break;
case MICROSECOND:
units = 1000000;
break;
case NANOSECOND:
units = 1000000000;
break;
default:
throw invalid("Unsupported timestamp unit");
}
long nanos = Math.floorMod(timestamp, units) * (1000000000 / units);
if (nanos % 1000 != 0) {
throw invalid("Timestamp parameter exceeds microsecond precision");
}
LocalDateTime time = LocalDateTime.ofEpochSecond(
Math.floorDiv(timestamp, units), (int) nanos, ZoneOffset.UTC);
checkYear(time.getYear());
if (type instanceof TimeStampTzType) {
// Zoned Arrow timestamps already encode a UTC instant. A DATETIMEV2 literal would
// reinterpret these fields in the session timezone when cast to TIMESTAMPTZ.
return new TimestampTzLiteral((TimeStampTzType) type, time.getYear(), time.getMonthValue(),
time.getDayOfMonth(), time.getHour(), time.getMinute(), time.getSecond(),
time.getNano() / 1000);
}
return new DateTimeV2Literal((DateTimeV2Type) type, time.getYear(), time.getMonthValue(),
time.getDayOfMonth(), time.getHour(), time.getMinute(), time.getSecond(),
time.getNano() / 1000);
}
}
private static void checkYear(int year) {
if (year < 0 || year > 9999) {
throw invalid("Date parameter is outside the supported year range 0000..9999");
}
}
private static FlightRuntimeException invalid(String message) {
return CallStatus.INVALID_ARGUMENT.withDescription(message).toRuntimeException();
}
}
@@ -27,6 +27,7 @@ import org.apache.doris.catalog.ScalarType;
import org.apache.doris.catalog.StructField;
import org.apache.doris.catalog.StructType;
import org.apache.doris.catalog.Type;
import org.apache.doris.common.Pair;
import org.apache.doris.datasource.CatalogIf;
import org.apache.doris.datasource.es.EsExternalCatalog;
import org.apache.doris.datasource.lance.LanceExternalCatalog;
@@ -38,7 +39,11 @@ import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import org.apache.doris.nereids.parser.NereidsParser;
import org.apache.doris.nereids.parser.SqlDialectHelper;
import org.apache.doris.nereids.rules.rewrite.CheckPrivileges;
import org.apache.doris.nereids.trees.expressions.Placeholder;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
import org.apache.doris.nereids.trees.plans.PlaceholderId;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.PrepareCommandPlanner;
import org.apache.doris.nereids.trees.plans.commands.AlterTableCommand;
@@ -68,6 +73,8 @@ import org.apache.doris.nereids.trees.plans.commands.insert.InsertOverwriteTable
import org.apache.doris.nereids.trees.plans.commands.merge.MergeIntoCommand;
import org.apache.doris.nereids.trees.plans.commands.use.SwitchCommand;
import org.apache.doris.nereids.trees.plans.commands.use.UseCommand;
import org.apache.doris.nereids.types.DataType;
import org.apache.doris.nereids.types.NullType;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.QueryState;
import org.apache.doris.qe.ResultSetMetaData;
@@ -97,6 +104,19 @@ final class FlightSqlQuerySchema {
}
static Schema analyze(ConnectContext context, String query) throws Exception {
return analyze(context, query, null);
}
static Schema analyze(ConnectContext context, String query, List<Literal> parameters) throws Exception {
return analyze(context, query, parameters, false).first;
}
static Pair<Schema, Schema> prepare(ConnectContext context, String query) throws Exception {
return analyze(context, query, null, true);
}
private static Pair<Schema, Schema> analyze(ConnectContext context, String query,
List<Literal> parameters, boolean preparing) throws Exception {
synchronized (context) {
ConnectContext previousThreadContext = ConnectContext.get();
StatementContext previousStatement = context.getStatementContext();
@@ -149,12 +169,33 @@ final class FlightSqlQuerySchema {
StatementContext statementContext = statement.getStatementContext();
context.setStatementContext(statementContext);
statementContext.setParsedStatement(statement);
if (!statementContext.getPlaceholders().isEmpty()) {
throw CallStatus.UNIMPLEMENTED.withDescription(
"Flight SQL parameter binding is not supported").toRuntimeException();
int parameterCount = statementContext.getPlaceholders().size();
if (parameterCount > FlightSqlParameters.MAX_PARAMETERS) {
throw CallStatus.INVALID_ARGUMENT.withDescription("Too many query parameters (maximum 1024)")
.toRuntimeException();
}
Plan plan = statement.getLogicalPlan();
Map<PlaceholderId, DataType> parameterTypes = new HashMap<>();
if (parameterCount > 0) {
if (statements.size() != 1 || plan instanceof Command) {
throw CallStatus.INVALID_ARGUMENT.withDescription(
"Parameters require a single query statement").toRuntimeException();
}
if (preparing) {
statementContext.setPrepareStage(true);
parameterTypes.putAll(FlightSqlParameters.inferCastTypes(plan));
parameters = new ArrayList<>();
for (Placeholder placeholder : statementContext.getPlaceholders()) {
// Unknown values must not acquire the synthetic STRING type used by MySQL Prepare.
parameters.add(new NullLiteral(parameterTypes.getOrDefault(
placeholder.getPlaceholderId(), NullType.INSTANCE)));
}
}
}
if (parameters != null || parameterCount > 0) {
FlightSqlParameters.bind(statementContext, parameters);
}
List<Field> fields = new ArrayList<>();
Plan plan = statement.getLogicalPlan();
if (plan instanceof Command) {
resolveNamespace(context, plan, scopedDatabases);
ResultSetMetaData metadata = commandMetadata(context, (Command) plan);
@@ -194,12 +235,26 @@ final class FlightSqlQuerySchema {
Plan analyzed = cascades.getRewritePlan();
// PrepareCommandPlanner stops before the rewrite phase that normally checks privileges.
new CheckPrivileges().rewriteRoot(analyzed, cascades.getCurrentJobContext());
if (preparing) {
statementContext.getIdToComparisonSlot().forEach((id, slot) ->
parameterTypes.putIfAbsent(id, slot.getDataType()));
}
for (Slot slot : analyzed.getOutput()) {
fields.add(field(slot.getName(), slot.getDataType().toCatalogDataType(), slot.nullable(),
true, context.getSessionVariable().getTimeZone()));
}
}
return new Schema(fields);
List<Field> parameterFields = new ArrayList<>();
for (int i = 0; preparing && i < parameterCount; i++) {
DataType type = parameterTypes.getOrDefault(
statementContext.getPlaceholders().get(i).getPlaceholderId(), NullType.INSTANCE);
Field parameterField = field(String.valueOf(i), type.toCatalogDataType(), true,
true, context.getSessionVariable().getTimeZone());
// Result schemas support more Arrow types than parameter uploads do.
parameterFields.add(FlightSqlParameters.supportsParameterType(parameterField.getType())
? parameterField : Field.nullable(String.valueOf(i), new ArrowType.Null()));
}
return Pair.of(new Schema(fields), new Schema(parameterFields));
} finally {
try {
List<AutoCloseable> resources = new ArrayList<>();
@@ -57,6 +57,9 @@ public class FlightSqlConnectPoolMgr extends ConnectPoolMgr {
@Override
public void unregisterConnection(ConnectContext ctx) {
// Use the short-lived prepared-state lock so teardown never waits for a running query
// holding the connection monitor before it can cancel that query.
ctx.closePreparedQueries();
// All Flight SQL session teardown paths (idle/query timeout, bearer token expiry, and
// explicit CloseSession) reach here. Release channel-cached Arrow results before removing
// the context from the pool.
@@ -182,11 +182,20 @@ public class DorisFlightSqlSchemaTest extends TestWithFeService {
@Test
void invalidSqlIsRejectedAndConnectionRecovers() throws Exception {
for (String sql : Arrays.asList("SELEC 1", "SELECT missing_column FROM schema_input")) {
Assertions.assertThrows(FlightRuntimeException.class, () -> schema(sql));
for (String sql : Arrays.asList("SELEC 1", "SELECT missing_column FROM schema_input",
"SELECT CAST(? AS BIGINT) FROM missing_schema_table")) {
FlightRuntimeException schemaFailure = Assertions.assertThrows(
FlightRuntimeException.class, () -> schema(sql));
Assertions.assertEquals(FlightStatusCode.INVALID_ARGUMENT, schemaFailure.status().code());
Mockito.clearInvocations(connectContext);
Assertions.assertThrows(java.util.concurrent.ExecutionException.class, () -> prepare(sql));
java.util.concurrent.ExecutionException prepareFailure = Assertions.assertThrows(
java.util.concurrent.ExecutionException.class, () -> prepare(sql));
Assertions.assertTrue(prepareFailure.getCause() instanceof FlightRuntimeException);
Assertions.assertEquals(FlightStatusCode.INVALID_ARGUMENT,
((FlightRuntimeException) prepareFailure.getCause()).status().code());
Mockito.verify(connectContext, Mockito.never()).addPreparedQuery(Mockito.anyString(), Mockito.anyString(), Mockito.any());
Mockito.verify(connectContext, Mockito.never()).addPreparedQuery(
Mockito.anyString(), Mockito.anyString(), Mockito.any(), Mockito.anyInt());
Assertions.assertEquals("id", preparedSchema("SELECT id FROM schema_input")
.getFields().get(0).getName());
}
@@ -0,0 +1,41 @@
// 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.doris.service.arrowflight;
import org.apache.doris.service.arrowflight.sessions.FlightSessionsManager;
import org.apache.arrow.flight.FlightProducer.CallContext;
import org.apache.arrow.flight.FlightProducer.StreamListener;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery;
import org.junit.Assert;
import org.junit.Test;
import org.mockito.Mockito;
public class FlightSqlParameterBindingTest {
@Test
public void acceptsParameterUpload() throws Exception {
try (DorisFlightSqlProducer producer = new DorisFlightSqlProducer(
Location.forGrpcInsecure("127.0.0.1", 0), Mockito.mock(FlightSessionsManager.class))) {
Assert.assertNotNull(producer.acceptPutPreparedStatementQuery(
CommandPreparedStatementQuery.getDefaultInstance(), Mockito.mock(CallContext.class),
Mockito.mock(FlightStream.class), Mockito.mock(StreamListener.class)));
}
}
}
@@ -0,0 +1,666 @@
// 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.doris.service.arrowflight;
import org.apache.doris.analysis.StatementBase;
import org.apache.doris.nereids.StatementContext;
import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import org.apache.doris.nereids.parser.NereidsParser;
import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral;
import org.apache.doris.nereids.trees.expressions.literal.DateTimeLiteral;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
import org.apache.doris.nereids.trees.expressions.literal.StringLiteral;
import org.apache.doris.nereids.trees.expressions.literal.TimestampTzLiteral;
import org.apache.doris.nereids.types.DateTimeV2Type;
import org.apache.doris.nereids.types.StringType;
import org.apache.doris.nereids.types.TimeStampTzType;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.StmtExecutor;
import org.apache.doris.service.arrowflight.sessions.FlightSessionsManager;
import org.apache.doris.service.arrowflight.sessions.FlightSqlConnectContext;
import org.apache.doris.utframe.TestWithFeService;
import com.google.common.base.Strings;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import org.apache.arrow.flight.Action;
import org.apache.arrow.flight.AsyncPutListener;
import org.apache.arrow.flight.CallHeaders;
import org.apache.arrow.flight.FlightClient;
import org.apache.arrow.flight.FlightDescriptor;
import org.apache.arrow.flight.FlightProducer.CallContext;
import org.apache.arrow.flight.FlightRuntimeException;
import org.apache.arrow.flight.FlightServer;
import org.apache.arrow.flight.FlightStatusCode;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.auth2.CallHeaderAuthenticator;
import org.apache.arrow.flight.sql.FlightSqlClient;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionClosePreparedStatementRequest;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest;
import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult;
import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.BigIntVector;
import org.apache.arrow.vector.DateDayVector;
import org.apache.arrow.vector.DecimalVector;
import org.apache.arrow.vector.Float8Vector;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.TimeStampMicroVector;
import org.apache.arrow.vector.TimeStampVector;
import org.apache.arrow.vector.VarBinaryVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.mockito.MockedConstruction;
import org.mockito.Mockito;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.math.BigDecimal;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.Types;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
public class FlightSqlPreparedQueryTest extends TestWithFeService {
private RootAllocator allocator;
private FlightServer server;
private FlightClient client;
private FlightSqlClient sqlClient;
private DorisFlightSqlProducer producer;
private FlightSqlConnectContext flightContext;
@Override
protected void runBeforeAll() throws Exception {
allocator = new RootAllocator();
flightContext = new FlightSqlConnectContext("parameter-peer");
flightContext.setEnv(connectContext.getEnv());
flightContext.setCurrentUserIdentity(connectContext.getCurrentUserIdentity());
flightContext.setSessionVariable(connectContext.getSessionVariable());
FlightSessionsManager sessions = Mockito.mock(FlightSessionsManager.class);
Mockito.when(sessions.getConnectContext("parameter-peer")).thenReturn(flightContext);
producer = new DorisFlightSqlProducer(Location.forGrpcInsecure("127.0.0.1", 0), sessions);
server = FlightServer.builder(allocator, Location.forGrpcInsecure("127.0.0.1", 0), producer)
.headerAuthenticator(headers -> new CallHeaderAuthenticator.AuthResult() {
@Override
public String getPeerIdentity() {
return "parameter-peer";
}
@Override
public void appendToOutgoingHeaders(CallHeaders headers) {
headers.insert("authorization", "Bearer parameter-peer");
}
}).build().start();
client = FlightClient.builder(allocator, server.getLocation()).build();
sqlClient = new FlightSqlClient(client);
}
@Override
protected void runAfterAll() throws Exception {
sqlClient.close();
server.close();
java.lang.reflect.Field executor = DorisFlightSqlProducer.class.getDeclaredField("executorService");
executor.setAccessible(true);
((ExecutorService) executor.get(producer)).shutdownNow();
producer.close();
flightContext.getFlightSqlChannel().close();
allocator.close();
}
private CommandPreparedStatementQuery prepare(String sql) throws Exception {
Action action = new Action("CreatePreparedStatement", Any.pack(
ActionCreatePreparedStatementRequest.newBuilder().setQuery(sql).build()).toByteArray());
ActionCreatePreparedStatementResult prepared = Any.parseFrom(client.doAction(action).next().getBody())
.unpack(ActionCreatePreparedStatementResult.class);
Assertions.assertFalse(prepared.getDatasetSchema().isEmpty());
return CommandPreparedStatementQuery.newBuilder()
.setPreparedStatementHandle(prepared.getPreparedStatementHandle()).build();
}
private void bind(CommandPreparedStatementQuery command, VectorSchemaRoot root) {
AsyncPutListener listener = new AsyncPutListener();
FlightClient.ClientStreamListener writer = client.startPut(
FlightDescriptor.command(Any.pack(command).toByteArray()), root, listener);
writer.putNext();
writer.completed();
listener.getResult();
}
private String id(CommandPreparedStatementQuery command) {
return command.getPreparedStatementHandle().toStringUtf8().substring("parameter-peer:".length());
}
private Schema schema(CommandPreparedStatementQuery command) {
return client.getSchema(FlightDescriptor.command(Any.pack(command).toByteArray())).getSchema();
}
@Test
public void jdbcPrepareAdvertisesResultAndParameterTypes() throws Exception {
Class.forName("org.apache.arrow.driver.jdbc.ArrowFlightJdbcDriver");
try (Connection connection = DriverManager.getConnection("jdbc:arrow-flight-sql://127.0.0.1:"
+ server.getPort() + "?useEncryption=false", "fixture", "fixture");
PreparedStatement statement = connection.prepareStatement(
"SELECT CAST(? AS BIGINT) AS n, CAST(? AS STRING) AS s, CAST(? AS DOUBLE) AS d")) {
Assertions.assertEquals(3, statement.getMetaData().getColumnCount());
for (int i = 0; i < 3; i++) {
int type = new int[] {Types.BIGINT, Types.VARCHAR, Types.DOUBLE}[i];
Assertions.assertEquals(type, statement.getMetaData().getColumnType(i + 1));
Assertions.assertEquals(type, statement.getParameterMetaData().getParameterType(i + 1));
}
}
}
@Test
public void jdbcPrepareInfersRangePredicateTypes() throws Exception {
Class.forName("org.apache.arrow.driver.jdbc.ArrowFlightJdbcDriver");
try (Connection connection = DriverManager.getConnection("jdbc:arrow-flight-sql://127.0.0.1:"
+ server.getPort() + "?useEncryption=false", "fixture", "fixture");
PreparedStatement statement = connection.prepareStatement(
"SELECT number FROM numbers(\"number\"=\"10\") WHERE number >= ? AND ? > number")) {
Assertions.assertEquals(Types.BIGINT, statement.getParameterMetaData().getParameterType(1));
Assertions.assertEquals(Types.BIGINT, statement.getParameterMetaData().getParameterType(2));
Assertions.assertEquals(Types.BIGINT, statement.getMetaData().getColumnType(1));
}
}
@Test
public void jdbcPreparePreservesInnerCastsAndSubqueryTypes() throws Exception {
Class.forName("org.apache.arrow.driver.jdbc.ArrowFlightJdbcDriver");
try (Connection connection = DriverManager.getConnection("jdbc:arrow-flight-sql://127.0.0.1:"
+ server.getPort() + "?useEncryption=false", "fixture", "fixture");
PreparedStatement statement = connection.prepareStatement(
"SELECT CAST(CAST(? AS BIGINT) AS STRING), (SELECT CAST(? AS DOUBLE))")) {
Assertions.assertEquals(Types.BIGINT, statement.getParameterMetaData().getParameterType(1));
Assertions.assertEquals(Types.DOUBLE, statement.getParameterMetaData().getParameterType(2));
Assertions.assertEquals(Types.VARCHAR, statement.getMetaData().getColumnType(1));
Assertions.assertEquals(Types.DOUBLE, statement.getMetaData().getColumnType(2));
}
}
@Test
public void unknownParameterTypesAreNotInvented() throws Exception {
try (FlightSqlClient.PreparedStatement statement = sqlClient.prepare("SELECT ? AS value")) {
Assertions.assertEquals(new ArrowType.Null(),
statement.getParameterSchema().getFields().get(0).getType());
Assertions.assertEquals(new ArrowType.Null(),
statement.getResultSetSchema().getFields().get(0).getType());
}
}
@Test
public void invalidSqlReportsClientErrorOverFlight() throws Exception {
for (String query : Arrays.asList("SELEC 1", "SELECT definitely_missing_schema_column",
"SELECT CAST(? AS BIGINT) FROM missing_parameter_table")) {
FlightRuntimeException prepareFailure = Assertions.assertThrows(FlightRuntimeException.class,
() -> sqlClient.prepare(query));
Assertions.assertEquals(FlightStatusCode.INVALID_ARGUMENT, prepareFailure.status().code());
FlightRuntimeException schemaFailure = Assertions.assertThrows(FlightRuntimeException.class,
() -> sqlClient.getExecuteSchema(query));
Assertions.assertEquals(FlightStatusCode.INVALID_ARGUMENT, schemaFailure.status().code());
try (FlightSqlClient.PreparedStatement statement = sqlClient.prepare("SELECT 1 AS id")) {
Assertions.assertEquals("id", statement.getResultSetSchema().getFields().get(0).getName());
}
}
}
@Test
public void unsupportedParameterVectorsDoNotRestrictResultTypes() throws Exception {
try (FlightSqlClient.PreparedStatement statement = sqlClient.prepare("SELECT CAST(? AS ARRAY<INT>)")) {
Assertions.assertEquals(new ArrowType.Null(),
statement.getParameterSchema().getFields().get(0).getType());
Assertions.assertEquals(new ArrowType.List(),
statement.getResultSetSchema().getFields().get(0).getType());
}
try (FlightSqlClient.PreparedStatement statement = sqlClient.prepare(
"SELECT CAST(CAST(? AS STRING) AS ARRAY<INT>)")) {
Assertions.assertEquals(new ArrowType.Utf8(),
statement.getParameterSchema().getFields().get(0).getType());
Assertions.assertEquals(new ArrowType.List(),
statement.getResultSetSchema().getFields().get(0).getType());
}
}
@Test
public void bindsAndRebindsIntegerQueryOverFlight() throws Exception {
CommandPreparedStatementQuery command = prepare("SELECT CAST(? AS BIGINT) AS value");
Assertions.assertEquals(new ArrowType.Int(64, true), schema(command).getFields().get(0).getType());
Assertions.assertThrows(FlightRuntimeException.class,
() -> client.getInfo(FlightDescriptor.command(Any.pack(command).toByteArray())));
try (BigIntVector value = new BigIntVector("parameter", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(value)) {
for (long expected : new long[] {42, -7, Long.MAX_VALUE}) {
value.setSafe(0, expected);
root.setRowCount(1);
bind(command, root);
Assertions.assertEquals(new ArrowType.Int(64, true), schema(command).getFields().get(0).getType());
Assertions.assertEquals(expected, flightContext.getPreparedQueryParameters(id(command)).get(0).getValue());
}
} finally {
flightContext.removePreparedQuery(id(command));
}
}
@Test
public void bindsMixedTypesAndPreservesQuotedText() throws Exception {
String query = "SELECT CAST(? AS BIGINT) AS n, CAST(? AS STRING) AS text, CAST(? AS DOUBLE) AS fraction";
CommandPreparedStatementQuery command = prepare(query);
try (BigIntVector n = new BigIntVector("n", allocator);
VarCharVector text = new VarCharVector("text", allocator);
Float8Vector fraction = new Float8Vector("fraction", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(n, text, fraction)) {
String expected = "中文 ' ? \\ ; SELECT 2";
n.setSafe(0, 17);
text.setSafe(0, expected.getBytes(StandardCharsets.UTF_8));
fraction.setSafe(0, 2.5);
root.setRowCount(1);
bind(command, root);
List<Literal> parameters = flightContext.getPreparedQueryParameters(id(command));
Assertions.assertEquals(17L, parameters.get(0).getValue());
Assertions.assertEquals(expected, parameters.get(1).getValue());
Assertions.assertEquals(2.5, parameters.get(2).getValue());
Assertions.assertEquals(3, schema(command).getFields().size());
// Exercise the execution parser hook without requiring a BE in the FE unit test fixture.
try (FlightSqlConnectProcessor processor = new FlightSqlConnectProcessor(flightContext)) {
java.lang.reflect.Field bindings = FlightSqlConnectProcessor.class.getDeclaredField("parameters");
bindings.setAccessible(true);
bindings.set(processor, parameters);
List<StatementBase> statements = processor.parseWithFallback(query, query,
flightContext.getSessionVariable());
StatementContext statement = ((LogicalPlanAdapter) statements.get(0))
.getStatementContext();
try {
for (int i = 0; i < parameters.size(); i++) {
Assertions.assertEquals(parameters.get(i), statement.getIdToPlaceholderRealExpr().get(
statement.getPlaceholders().get(i).getPlaceholderId()));
}
} finally {
statement.close();
}
}
} finally {
flightContext.removePreparedQuery(id(command));
}
}
@Test
public void failedRebindInvalidatesOldValuesAndConnectionRecovers() throws Exception {
CommandPreparedStatementQuery command = prepare("SELECT ? AS value");
try (IntVector value = new IntVector("value", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(value)) {
value.setSafe(0, 1);
root.setRowCount(1);
bind(command, root);
value.setSafe(1, 2);
root.setRowCount(2);
FlightRuntimeException failure = Assertions.assertThrows(FlightRuntimeException.class,
() -> bind(command, root));
Assertions.assertEquals(FlightStatusCode.UNIMPLEMENTED, failure.status().code());
Assertions.assertNull(flightContext.getPreparedQueryParameters(id(command)));
Assertions.assertThrows(FlightRuntimeException.class,
() -> client.getInfo(FlightDescriptor.command(Any.pack(command).toByteArray())));
root.setRowCount(1);
value.setSafe(0, 3);
bind(command, root);
Assertions.assertEquals(new ArrowType.Int(32, true), schema(command).getFields().get(0).getType());
value.setNull(0);
bind(command, root);
Assertions.assertTrue(flightContext.getPreparedQueryParameters(id(command)).get(0) instanceof NullLiteral);
} finally {
flightContext.removePreparedQuery(id(command));
}
Assertions.assertNotNull(sqlClient.getExecuteSchema("SELECT 1"));
}
@Test
public void convertsDetachedValuesAndTypedNulls() {
try (IntVector value = new IntVector("value", allocator);
VarCharVector text = new VarCharVector("text", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(value, text)) {
value.setSafe(0, 12);
text.setNull(0);
root.setRowCount(1);
List<Literal> literals = FlightSqlParameters.convert(root);
value.setSafe(0, 99);
Assertions.assertEquals(12, literals.get(0).getValue());
Assertions.assertTrue(literals.get(1) instanceof NullLiteral);
Assertions.assertEquals(StringType.INSTANCE, literals.get(1).getDataType());
}
}
@Test
public void validatesWholeUploadInsteadOfUsingFirstBatch() {
try (IntVector value = new IntVector("value", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(value)) {
value.setSafe(0, 1);
root.setRowCount(1);
FlightStream stream = Mockito.mock(FlightStream.class);
Mockito.when(stream.next()).thenReturn(true, true, false);
Mockito.when(stream.getRoot()).thenReturn(root);
Assertions.assertThrows(Exception.class, () -> FlightSqlParameters.read(stream, 1));
Mockito.verify(stream, Mockito.times(3)).next();
FlightStream empty = Mockito.mock(FlightStream.class);
Assertions.assertThrows(Exception.class, () -> FlightSqlParameters.read(empty, 1));
Assertions.assertEquals(Arrays.asList(), FlightSqlParameters.read(empty, 0));
}
}
@Test
public void limitsRetainedParametersAcrossHandlesAndReleasesTheirBudget() {
ConnectContext connection = new ConnectContext();
// Test sources also target Java 8, where String.repeat is unavailable.
List<Literal> parameters = Arrays.asList(
new StringLiteral(Strings.repeat("x", 512 * 1024)));
for (int i = 0; i < 15; i++) {
connection.addPreparedQuery("p" + i, "SELECT ?", null, 1);
connection.setPreparedQueryParameters("p" + i, parameters, null);
}
connection.addPreparedQuery("overflow", "SELECT ?", null, 1);
Assertions.assertThrows(IllegalArgumentException.class,
() -> connection.setPreparedQueryParameters("overflow", parameters, null));
connection.removePreparedQuery("p0");
connection.setPreparedQueryParameters("overflow", parameters, null);
connection.beginPreparedQueryBinding("overflow");
connection.addPreparedQuery("replacement", "SELECT ?", null, 1);
connection.setPreparedQueryParameters("replacement", parameters, null);
long upload = connection.beginPreparedQueryBinding("replacement");
connection.clearPreparedQueries();
Assertions.assertFalse(connection.isPreparedQueryBindingCurrent("replacement", upload));
Assertions.assertNull(connection.getPreparedQuery("replacement"));
connection.addPreparedQuery("fresh", "SELECT ?", null, 1);
connection.setPreparedQueryParameters("fresh", parameters, null);
}
@Test
public void rejectsInvalidValuesAndUnsupportedTypes() {
try (VarCharVector text = new VarCharVector("text", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(text)) {
text.setSafe(0, new byte[] {(byte) 0xc3, (byte) 0x28});
root.setRowCount(1);
Assertions.assertThrows(FlightRuntimeException.class, () -> FlightSqlParameters.convert(root));
}
try (Float8Vector number = new Float8Vector("number", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(number)) {
number.setSafe(0, Double.POSITIVE_INFINITY);
root.setRowCount(1);
Assertions.assertThrows(FlightRuntimeException.class, () -> FlightSqlParameters.convert(root));
}
try (VarBinaryVector binary =
new VarBinaryVector("binary", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(binary)) {
binary.setSafe(0, new byte[] {0, 1});
root.setRowCount(1);
FlightRuntimeException error = Assertions.assertThrows(FlightRuntimeException.class,
() -> FlightSqlParameters.convert(root));
Assertions.assertEquals(FlightStatusCode.UNIMPLEMENTED, error.status().code());
}
}
@Test
public void rejectsForwardingInsteadOfLosingTypedBindings() throws Exception {
flightContext.setThreadLocalInfo();
LogicalPlanAdapter statement = (LogicalPlanAdapter) new NereidsParser()
.parseSQL("SELECT ?", flightContext.getSessionVariable()).get(0);
try {
FlightSqlParameters.bind(statement.getStatementContext(), Arrays.asList(
new BigIntLiteral(7)));
StmtExecutor executor = new StmtExecutor(flightContext, statement);
Method forward = StmtExecutor.class
.getDeclaredMethod("forwardToMaster");
forward.setAccessible(true);
InvocationTargetException failure = Assertions.assertThrows(
InvocationTargetException.class, () -> forward.invoke(executor));
Assertions.assertTrue(failure.getCause().getMessage().contains("connect to master FE"));
} finally {
statement.getStatementContext().close();
ConnectContext.remove();
}
}
@Test
public void preservesDecimalAndTemporalPrecision() {
try (DecimalVector decimal =
new DecimalVector("decimal", allocator, 20, 6);
DateDayVector date = new DateDayVector("date", allocator);
TimeStampMicroVector timestamp =
new TimeStampMicroVector("timestamp", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(decimal, date, timestamp)) {
BigDecimal expected = new BigDecimal("12345678901234.567890");
decimal.setSafe(0, expected);
date.setSafe(0, -1);
timestamp.setSafe(0, -1);
root.setRowCount(1);
List<Literal> values = FlightSqlParameters.convert(root);
Assertions.assertEquals(expected, values.get(0).getValue());
Assertions.assertEquals("1969-12-31", values.get(1).getStringValue());
Assertions.assertEquals("1969-12-31 23:59:59.999999", values.get(2).getStringValue());
date.setSafe(0, Integer.MAX_VALUE);
Assertions.assertThrows(FlightRuntimeException.class, () -> FlightSqlParameters.convert(root));
}
}
@Test
public void preservesTimestampSemanticsAndNulls() {
for (String zone : Arrays.asList(null, "", "UTC", "Asia/Shanghai", "America/New_York")) {
for (org.apache.arrow.vector.types.TimeUnit unit : org.apache.arrow.vector.types.TimeUnit.values()) {
ArrowType.Timestamp type = new ArrowType.Timestamp(unit, zone);
long nanosPerUnit = java.util.concurrent.TimeUnit.valueOf(unit.name() + "S").toNanos(1);
try (VectorSchemaRoot root = VectorSchemaRoot.create(
new Schema(Arrays.asList(Field.nullable("value", type))), allocator)) {
root.allocateNew();
TimeStampVector vector = (TimeStampVector) root.getVector(0);
for (long value : new long[] {-1000, 0, 1000}) {
vector.setSafe(0, value);
root.setRowCount(1);
Literal literal = FlightSqlParameters.convert(root).get(0);
Assertions.assertEquals(zone != null && !zone.isEmpty(), literal instanceof TimestampTzLiteral);
Assertions.assertEquals(java.time.LocalDateTime.of(1970, 1, 1, 0, 0)
.plusNanos(value * nanosPerUnit), ((DateTimeLiteral) literal).toJavaDateType());
}
vector.setNull(0);
Literal literal = FlightSqlParameters.convert(root).get(0);
Assertions.assertTrue(literal instanceof NullLiteral);
Assertions.assertEquals(zone != null && !zone.isEmpty()
? TimeStampTzType.of(6) : DateTimeV2Type.of(6), literal.getDataType());
}
}
}
}
@Test
public void rejectsTimezoneTimestampPrecisionLossAndOverflow() {
for (org.apache.arrow.vector.types.TimeUnit unit : Arrays.asList(
org.apache.arrow.vector.types.TimeUnit.NANOSECOND, org.apache.arrow.vector.types.TimeUnit.SECOND)) {
try (VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(Arrays.asList(
Field.nullable("value", new ArrowType.Timestamp(unit, "UTC")))), allocator)) {
root.allocateNew();
TimeStampVector vector = (TimeStampVector) root.getVector(0);
long[] values = unit == org.apache.arrow.vector.types.TimeUnit.NANOSECOND
? new long[] {-1, 1, 1001} : new long[] {Long.MIN_VALUE, Long.MAX_VALUE};
for (long value : values) {
vector.setSafe(0, value);
root.setRowCount(1);
FlightStream stream = Mockito.mock(FlightStream.class);
Mockito.when(stream.next()).thenReturn(true, false);
Mockito.when(stream.getRoot()).thenReturn(root);
FlightRuntimeException failure = Assertions.assertThrows(FlightRuntimeException.class,
() -> FlightSqlParameters.read(stream, 1));
Assertions.assertEquals(FlightStatusCode.INVALID_ARGUMENT, failure.status().code());
}
}
}
}
@Test
public void bindsTimezoneTimestampInNonUtcSession() throws Exception {
String previousZone = flightContext.getSessionVariable().getTimeZone();
try {
for (String zone : Arrays.asList("Asia/Shanghai", "America/New_York")) {
flightContext.getSessionVariable().setTimeZone(zone);
try (FlightSqlClient.PreparedStatement statement = sqlClient.prepare("SELECT CAST(? AS TIMESTAMPTZ(6))")) {
Assertions.assertEquals(new ArrowType.Timestamp(
org.apache.arrow.vector.types.TimeUnit.MICROSECOND, zone),
statement.getParameterSchema().getFields().get(0).getType());
}
CommandPreparedStatementQuery command = prepare("SELECT CAST(? AS TIMESTAMPTZ(6))");
try (VectorSchemaRoot root = VectorSchemaRoot.create(new Schema(Arrays.asList(Field.nullable("value",
new ArrowType.Timestamp(org.apache.arrow.vector.types.TimeUnit.MICROSECOND, "UTC")))), allocator)) {
root.allocateNew();
TimeStampVector vector = (TimeStampVector) root.getVector(0);
for (long value : new long[] {0, -1, 1234567}) {
vector.setSafe(0, value);
root.setRowCount(1);
bind(command, root);
TimestampTzLiteral literal = (TimestampTzLiteral) flightContext
.getPreparedQueryParameters(id(command)).get(0);
Assertions.assertEquals(java.time.LocalDateTime.of(1970, 1, 1, 0, 0)
.plusNanos(value * 1000), literal.toJavaDateType());
Assertions.assertEquals(new ArrowType.Timestamp(
org.apache.arrow.vector.types.TimeUnit.MICROSECOND, zone),
schema(command).getFields().get(0).getType());
}
vector.setNull(0);
bind(command, root);
Assertions.assertEquals(TimeStampTzType.of(6), flightContext
.getPreparedQueryParameters(id(command)).get(0).getDataType());
} finally {
client.doAction(new Action("ClosePreparedStatement", Any.pack(
ActionClosePreparedStatementRequest.newBuilder()
.setPreparedStatementHandle(command.getPreparedStatementHandle()).build())
.toByteArray())).forEachRemaining(result -> { });
}
}
} finally {
flightContext.getSessionVariable().setTimeZone(previousZone);
}
}
@Test
public void rejectsForeignAndClosedHandles() throws Exception {
CommandPreparedStatementQuery command = prepare("SELECT ?");
CommandPreparedStatementQuery foreign = CommandPreparedStatementQuery.newBuilder()
.setPreparedStatementHandle(ByteString.copyFromUtf8("another-peer:handle")).build();
FlightRuntimeException failure = Assertions.assertThrows(FlightRuntimeException.class, () -> schema(foreign));
Assertions.assertEquals(FlightStatusCode.INVALID_ARGUMENT, failure.status().code());
ActionClosePreparedStatementRequest close =
ActionClosePreparedStatementRequest.newBuilder()
.setPreparedStatementHandle(command.getPreparedStatementHandle()).build();
client.doAction(new Action("ClosePreparedStatement", Any.pack(close).toByteArray())).forEachRemaining(r -> { });
FlightRuntimeException closed = Assertions.assertThrows(FlightRuntimeException.class, () -> schema(command));
Assertions.assertEquals(FlightStatusCode.NOT_FOUND, closed.status().code());
}
@Test
public void bindsRangePredicatesAndRejectsWrongParameterCount() throws Exception {
CommandPreparedStatementQuery command = prepare("SELECT number FROM numbers(\"number\"=\"10\") "
+ "WHERE number >= ? AND number < ? ORDER BY number");
try (BigIntVector lower = new BigIntVector("lower", allocator);
BigIntVector upper = new BigIntVector("upper", allocator);
VectorSchemaRoot root = VectorSchemaRoot.of(lower, upper)) {
lower.setSafe(0, 2);
upper.setSafe(0, 5);
root.setRowCount(1);
bind(command, root);
Assertions.assertEquals(new ArrowType.Int(64, true), schema(command).getFields().get(0).getType());
}
try (IntVector only = new IntVector("only", allocator);
VectorSchemaRoot wrong = VectorSchemaRoot.of(only)) {
only.setSafe(0, 1);
wrong.setRowCount(1);
FlightRuntimeException error = Assertions.assertThrows(FlightRuntimeException.class,
() -> bind(command, wrong));
Assertions.assertEquals(FlightStatusCode.INVALID_ARGUMENT, error.status().code());
Assertions.assertNull(flightContext.getPreparedQueryParameters(id(command)));
} finally {
flightContext.removePreparedQuery(id(command));
}
}
@Test
public void sessionCloseDoesNotWaitForExecutionMonitor() throws Exception {
ConnectContext connection = new ConnectContext();
connection.addPreparedQuery("p", "SELECT ?", null, 1);
long upload = connection.beginPreparedQueryBinding("p");
ExecutorService executor = Executors.newSingleThreadExecutor();
try {
synchronized (connection) {
executor.submit(connection::closePreparedQueries).get(5, TimeUnit.SECONDS);
}
Assertions.assertFalse(connection.isPreparedQueryBindingCurrent("p", upload));
Assertions.assertThrows(IllegalStateException.class,
() -> connection.addPreparedQuery("new", "SELECT ?", null, 1));
} finally {
executor.shutdownNow();
}
}
@Test
public void closedHandleCannotCommitAnUploadOrLeaveUnreachableResults() throws Exception {
ConnectContext connection = new ConnectContext();
connection.addPreparedQuery("p", "SELECT ?", null, 1);
long version = connection.beginPreparedQueryBinding("p");
connection.closePreparedQueries();
Assertions.assertFalse(connection.setPreparedQueryParameters("p", Arrays.asList(new BigIntLiteral(1)),
null, version));
Assertions.assertEquals(-1, connection.getPreparedQueryParameterCount("p"));
Assertions.assertNull(connection.getPreparedQueryParameters("p"));
Assertions.assertEquals(-1, connection.beginPreparedQueryBinding("p"));
String query = "SET query_timeout = 300";
flightContext.getState().reset();
flightContext.addPreparedQuery("close-race", query, FlightSqlQuerySchema.analyze(flightContext, query));
CallContext call = Mockito.mock(CallContext.class);
Mockito.when(call.peerIdentity()).thenReturn("parameter-peer");
CommandPreparedStatementQuery command = CommandPreparedStatementQuery.newBuilder()
.setPreparedStatementHandle(ByteString.copyFromUtf8("parameter-peer:close-race")).build();
StmtExecutor deferred = Mockito.mock(StmtExecutor.class);
boolean previousLocal = flightContext.isReturnResultFromLocal();
try (MockedConstruction<FlightSqlConnectProcessor> ignored = Mockito.mockConstruction(
FlightSqlConnectProcessor.class, (processor, context) -> {
Mockito.doAnswer(invocation -> {
flightContext.clearPreparedQueries();
flightContext.setReturnResultFromLocal(true);
flightContext.addFlightSqlDeferredExecutor(deferred);
return null;
}).when(processor).handleQuery(Mockito.anyString());
})) {
FlightRuntimeException error = Assertions.assertThrows(FlightRuntimeException.class,
() -> producer.getFlightInfoPreparedStatement(
command, call, FlightDescriptor.command(Any.pack(command).toByteArray())));
Assertions.assertEquals(FlightStatusCode.NOT_FOUND, error.status().code());
Assertions.assertEquals(0, flightContext.getFlightSqlChannel().resultNum());
Mockito.verify(deferred).finalizeArrowFlightQuery();
} finally {
flightContext.setReturnResultFromLocal(previousLocal);
flightContext.closeFlightSqlDeferredExecutors();
flightContext.getFlightSqlChannel().reset();
}
}
}
@@ -0,0 +1,223 @@
// 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.
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CallOption
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CallOptions
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CloseSessionRequest
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightClient
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightRuntimeException
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightStatusCode
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.Location
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.sql.FlightSqlClient
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.memory.RootAllocator
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.BigIntVector
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.Float8Vector
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.TimeStampMicroTZVector
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.VarCharVector
import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.vector.VectorSchemaRoot
import java.sql.DriverManager
import java.sql.Types
suite("test_prepared_query_parameters", "arrow_flight_sql") {
def config = context.config.otherConfigs
def location = Location.forGrpcInsecure(config.get("extArrowFlightSqlHost"),
Integer.parseInt(config.get("extArrowFlightSqlPort")))
Class.forName("org.apache.arrow.driver.jdbc.ArrowFlightJdbcDriver")
def jdbcUrl = "jdbc:arrow-flight-sql://${config.get('extArrowFlightSqlHost')}:" +
"${config.get('extArrowFlightSqlPort')}?useEncryption=false"
DriverManager.getConnection(jdbcUrl, config.get("extArrowFlightSqlUser"),
config.get("extArrowFlightSqlPassword")).withCloseable { connection ->
// Exercise JDBC's schema-driven binder and query classifier, not only direct vector uploads.
connection.prepareStatement("SELECT CAST(? AS BIGINT) AS value").withCloseable { statement ->
assertEquals(1, statement.getMetaData().getColumnCount())
assertEquals(Types.BIGINT, statement.getParameterMetaData().getParameterType(1))
[42L, -7L, Long.MAX_VALUE].each { value ->
statement.setLong(1, value)
assertTrue(statement.execute())
statement.getResultSet().withCloseable { result ->
assertTrue(result.next())
assertEquals(value, result.getLong(1))
assertTrue(!result.next())
}
}
statement.setNull(1, Types.BIGINT)
statement.executeQuery().withCloseable { result ->
assertTrue(result.next())
assertEquals(null, result.getObject(1))
assertTrue(!result.next())
}
}
connection.prepareStatement("SELECT CAST(? AS BIGINT), CAST(? AS STRING), CAST(? AS DOUBLE)")
.withCloseable { statement ->
statement.setLong(1, 17L)
statement.setString(2, "quoted ' ? 中文")
statement.setDouble(3, 2.5d)
statement.executeQuery().withCloseable { result ->
assertTrue(result.next())
assertEquals(17L, result.getLong(1))
assertEquals("quoted ' ? 中文", result.getString(2))
assertEquals(2.5d, result.getDouble(3))
assertTrue(!result.next())
}
}
connection.prepareStatement('SELECT number FROM numbers("number"="10") '
+ 'WHERE number >= ? AND ? > number ORDER BY number').withCloseable { statement ->
[[2L, 5L], [6L, 9L]].each { bounds ->
statement.setLong(1, bounds[0])
statement.setLong(2, bounds[1])
statement.executeQuery().withCloseable { result ->
def rows = []
while (result.next()) {
rows.add(result.getLong(1))
}
assertEquals((bounds[0]..<bounds[1]).toList(), rows)
}
}
}
}
new RootAllocator().withCloseable { allocator ->
FlightClient.builder(allocator, location).build().withCloseable { flight ->
def token = flight.authenticateBasicToken(config.get("extArrowFlightSqlUser"),
config.get("extArrowFlightSqlPassword")).get()
CallOption[] options = [token, CallOptions.timeout(30, java.util.concurrent.TimeUnit.SECONDS)]
def client = new FlightSqlClient(flight)
def fetch = { info ->
def rows = []
info.getEndpoints().each { endpoint ->
def resultLocation = endpoint.getLocations().isEmpty() ? location : endpoint.getLocations()[0]
FlightClient.builder(allocator, resultLocation).build().withCloseable { resultClient ->
resultClient.getStream(endpoint.getTicket(), options).withCloseable { stream ->
while (stream.next()) {
def root = stream.getRoot()
for (int row = 0; row < root.getRowCount(); row++) {
rows.add(root.getFieldVectors().collect { vector ->
def value = vector.getObject(row)
value == null ? null : value.toString()
})
}
}
}
}
}
rows
}
try {
def prepared = client.prepare("SELECT CAST(? AS BIGINT) AS value", options)
try {
new BigIntVector("value", allocator).withCloseable { value ->
VectorSchemaRoot.of(value).withCloseable { root ->
prepared.setParameters(root)
[42L, -7L, Long.MAX_VALUE].each { expected ->
value.setSafe(0, expected)
root.setRowCount(1)
assertEquals([[expected.toString()]], fetch(prepared.execute(options)))
}
value.setNull(0)
assertEquals([[null]], fetch(prepared.execute(options)))
value.setSafe(0, 1L)
value.setSafe(1, 2L)
root.setRowCount(2)
try {
prepared.execute(options)
assertTrue(false, "Multiple parameter rows must not silently execute only the first")
} catch (FlightRuntimeException e) {
assertEquals(FlightStatusCode.UNIMPLEMENTED, e.status().code())
}
root.setRowCount(1)
value.setSafe(0, 3L)
assertEquals([["3"]], fetch(prepared.execute(options)))
}
}
} finally {
// Close must use the same authenticated session that owns the prepared handle.
prepared.close(options)
}
prepared = client.prepare("SELECT CAST(? AS BIGINT), CAST(? AS STRING), CAST(? AS DOUBLE)", options)
try {
VectorSchemaRoot.of(new BigIntVector("n", allocator), new VarCharVector("s", allocator),
new Float8Vector("f", allocator)).withCloseable { root ->
String text = "quoted ' text ? with Unicode 中文"
root.getVector(0).setSafe(0, 17L)
root.getVector(1).setSafe(0, text.getBytes("UTF-8"))
root.getVector(2).setSafe(0, 2.5d)
root.setRowCount(1)
prepared.setParameters(root)
assertEquals([["17", text, "2.5"]], fetch(prepared.execute(options)))
}
} finally {
prepared.close(options)
}
prepared = client.prepare('SELECT number FROM numbers("number"="10") '
+ 'WHERE number >= ? AND number < ? ORDER BY number', options)
try {
VectorSchemaRoot.of(new BigIntVector("lower", allocator),
new BigIntVector("upper", allocator)).withCloseable { root ->
root.setRowCount(1)
prepared.setParameters(root)
[[2L, 5L], [6L, 9L]].each { bounds ->
root.getVector(0).setSafe(0, bounds[0])
root.getVector(1).setSafe(0, bounds[1])
root.setRowCount(1)
assertEquals((bounds[0]..<bounds[1]).collect { [it.toString()] },
fetch(prepared.execute(options)))
}
}
} finally {
prepared.close(options)
}
def originalZone = fetch(client.execute("SHOW VARIABLES LIKE 'time_zone'", options))[0][1]
try {
["UTC", "Asia/Shanghai", "America/New_York"].each { sessionZone ->
fetch(client.execute("SET time_zone='${sessionZone}'".toString(), options))
prepared = client.prepare('SELECT number FROM numbers("number"="3") '
+ 'WHERE CAST(? AS TIMESTAMPTZ(6)) = CAST(? AS TIMESTAMPTZ(6)) ORDER BY number', options)
try {
assertEquals(sessionZone,
prepared.getParameterSchema().getFields()[0].getType().getTimezone())
["UTC", "America/New_York"].each { inputZone ->
VectorSchemaRoot.of(new TimeStampMicroTZVector("instant", allocator, inputZone),
new VarCharVector("expected", allocator)).withCloseable { root ->
prepared.setParameters(root)
// The timezone annotation must not shift the epoch or change predicate matches.
[[0L, "1970-01-01 00:00:00+00:00"],
[-1L, "1969-12-31 23:59:59.999999+00:00"],
[1234567L, "1970-01-01 00:00:01.234567+00:00"],
[1710054000000000L, "2024-03-10 07:00:00+00:00"]].each { value ->
root.getVector(0).setSafe(0, value[0])
root.getVector(1).setSafe(0, value[1].getBytes("UTF-8"))
root.setRowCount(1)
assertEquals([["0"], ["1"], ["2"]], fetch(prepared.execute(options)))
}
root.getVector(0).setNull(0)
assertEquals([], fetch(prepared.execute(options)))
}
}
} finally {
prepared.close(options)
}
}
} finally {
fetch(client.execute("SET time_zone='${originalZone}'".toString(), options))
}
assertEquals([["1"]], fetch(client.execute("SELECT 1", options)))
} finally {
flight.closeSession(new CloseSessionRequest(), options)
}
}
}
}