This is an automated email from the ASF dual-hosted git repository. yiguolei pushed a commit to branch branch-4.2 in repository https://gitbox.apache.org/repos/asf/doris.git
commit cd3af8974c7a564f489a60b82137ef3bc13f9faa Author: Gabriel <[email protected]> AuthorDate: Fri Oct 9 14:37:12 2026 +0800 [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 --- .../nereids/rules/analysis/ExpressionAnalyzer.java | 5 +- .../java/org/apache/doris/qe/ConnectContext.java | 176 +++++- .../java/org/apache/doris/qe/StmtExecutor.java | 6 + .../arrowflight/DorisFlightSqlProducer.java | 132 +++- .../arrowflight/FlightSqlConnectProcessor.java | 38 ++ .../service/arrowflight/FlightSqlParameters.java | 336 +++++++++++ .../service/arrowflight/FlightSqlQuerySchema.java | 65 +- .../sessions/FlightSqlConnectPoolMgr.java | 3 + .../arrowflight/DorisFlightSqlSchemaTest.java | 15 +- .../arrowflight/FlightSqlParameterBindingTest.java | 41 ++ .../arrowflight/FlightSqlPreparedQueryTest.java | 666 +++++++++++++++++++++ .../test_prepared_query_parameters.groovy | 223 +++++++ 12 files changed, 1648 insertions(+), 58 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java index d9cc562d305..a68d539d284 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ExpressionAnalyzer.java @@ -816,7 +816,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(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java index 674c2a99548..e51a0031676 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java @@ -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 void addPreparedQuery(String preparedStatementId, String preparedQuery) { + synchronized (preparedQuerys) { + addPreparedQuery(preparedStatementId, preparedQuery, null); + } } - public synchronized void addPreparedQuery(String preparedStatementId, String preparedQuery, Schema schema) { - preparedQuerys.put(preparedStatementId, - new PreparedQuery(preparedQuery, getDefaultCatalog(), getDatabase(), schema)); + 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 synchronized Schema getPreparedQuerySchema(String preparedStatementId) { - PreparedQuery query = preparedQuerys.get(preparedStatementId); - return query == null ? null : query.schema; + public void addPreparedQuery(String id, String sql, Schema schema, int parameterCount) { + synchronized (preparedQuerys) { + addPreparedQuery(id, sql, schema); + preparedQuerys.get(id).parameterCount = parameterCount; + } } - public synchronized String getPreparedQuery(String preparedStatementId) { - PreparedQuery query = preparedQuerys.get(preparedStatementId); - if (query == null) { - return null; + public int getPreparedQueryParameterCount(String id) { + synchronized (preparedQuerys) { + PreparedQuery query = preparedQuerys.get(id); + return query == null ? -1 : query.parameterCount; } - // 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; + } + + 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; } - return query.sql; } - public synchronized void removePreparedQuery(String preparedStatementId) { - preparedQuerys.remove(preparedStatementId); + 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) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index 1919141ec00..3fd32020d92 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java @@ -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()); diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java index 2aae87d527b..13b3f506fb9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java @@ -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 diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java index 5942424750c..87f00e763b8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java @@ -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); diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlParameters.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlParameters.java new file mode 100644 index 00000000000..59d21fc71ce --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlParameters.java @@ -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(); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java index 678cab055a8..0ed2b77e92f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java +++ b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlQuerySchema.java @@ -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(); } - List<Field> fields = new ArrayList<>(); 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<>(); 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<>(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgr.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgr.java index c8854507e00..63d01a87a8a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgr.java +++ b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgr.java @@ -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. diff --git a/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlSchemaTest.java b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlSchemaTest.java index 70fb507e502..86fea280901 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlSchemaTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlSchemaTest.java @@ -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()); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlParameterBindingTest.java b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlParameterBindingTest.java new file mode 100644 index 00000000000..1f6c7cc2a7f --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlParameterBindingTest.java @@ -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))); + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlPreparedQueryTest.java b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlPreparedQueryTest.java new file mode 100644 index 00000000000..aa6d37582bb --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/FlightSqlPreparedQueryTest.java @@ -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(); + } + } + +} diff --git a/regression-test/suites/arrow_flight_sql_p0/test_prepared_query_parameters.groovy b/regression-test/suites/arrow_flight_sql_p0/test_prepared_query_parameters.groovy new file mode 100644 index 00000000000..e41a3bbbbaa --- /dev/null +++ b/regression-test/suites/arrow_flight_sql_p0/test_prepared_query_parameters.groovy @@ -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) + } + } + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
