diff --git a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlProducer.java b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlProducer.java index c498143a33981c..1bf4bdfa2945d7 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlProducer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlProducer.java @@ -197,13 +197,11 @@ private FlightInfo executeQueryStatement(String peerIdentity, ConnectContext con try { Preconditions.checkState(null != connectContext); Preconditions.checkState(!query.isEmpty()); - // Finalize the previous query's coordinator on this connection whose close was - // deferred (Arrow Flight keeps it alive across GetFlightInfo -> DoGet so the BE can - // fetch external-table splits during DoGet). By now the previous DoGet is done. #62259 - connectContext.closeFlightSqlDeferredExecutors(); - // After the previous query was executed, there was no getStreamStatement to take away the result. - connectContext.getFlightSqlChannel().reset(); - connectContext.clearFlightSqlEndpointsLocations(); + // Drops what the previous request left on the session: its deferred coordinator (Arrow + // Flight keeps it alive across GetFlightInfo -> DoGet so the BE can fetch external-table + // splits during DoGet, and by now that DoGet is done, #62259), a result no + // getStreamStatement took away, and its endpoints. + FlightProtocolAdapter.of(connectContext).beginRequest(); try (FlightSqlConnectProcessor flightSQLConnectProcessor = new FlightSqlConnectProcessor(connectContext)) { flightSQLConnectProcessor.handleQuery(query); if (connectContext.getState().getStateType() == MysqlStateType.ERR) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlConnectProcessor.java index b9a95ed77d1b32..0203e29a802e4c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlConnectProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlConnectProcessor.java @@ -64,7 +64,6 @@ public class FlightSqlConnectProcessor extends ConnectProcessor implements AutoC public FlightSqlConnectProcessor(ConnectContext context) { super(context); context.setThreadLocalInfo(); - context.setReturnResultFromLocal(true); } public Schema getArrowSchema() { diff --git a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java index 273a2f5bef71ff..fcda6a9aa82d45 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java @@ -30,6 +30,7 @@ import org.apache.doris.qe.ShowResultSet; import org.apache.doris.qe.StmtExecutor; import org.apache.doris.qe.protocol.ProtocolAdapter; +import org.apache.doris.thrift.TMasterOpRequest; import org.apache.doris.thrift.TResultSinkType; import com.google.common.annotations.VisibleForTesting; @@ -66,6 +67,14 @@ public class FlightProtocolAdapter implements ProtocolAdapter { private final Map preparedQuerys = new HashMap<>(); private String runningQuery; private final List endpointsLocations = Lists.newArrayList(); + // How many of endpointsLocations were registered before the statement being executed + // started: what an attempt of that statement registers comes after them, and only that is + // withdrawn when the statement is attempted again (beforeAttempt). + private int endpointsBeforeStatement = 0; + // Whether the result of the statement being executed is on this frontend (a SHOW, a SET, an + // EXPLAIN: cached on the channel for the client's DoGet) or on the backends the coordinator + // ran the query on, registered in endpointsLocations for the client to pull from. Set by the + // statement lifecycle hooks below. private boolean returnResultFromLocal = true; // Executors of already-planned queries whose results are produced on the BE and pulled later // during the DoGet phase. Their coordinators must stay alive until the BE finishes scanning: @@ -129,6 +138,100 @@ public boolean supportsSqlCacheReplay() { return false; } + /** + * The master returns a query result as MySQL wire packets, which cannot be turned into the + * Arrow batches a Flight client needs. The executor refuses to forward a query rather than + * let the master build a result set this frontend would discard and answer the client with a + * synthesized empty success. + */ + @Override + public boolean canReplayForwardedQueryResult() { + return false; + } + + /** + * A result this frontend materializes is cached with every column as a Utf8 vector, whatever + * its type ({@link FlightResultSender}). That is acceptable for the text a SHOW or an EXPLAIN + * produces, not for a SELECT a client expects typed Arrow data from, so a query the planner + * could answer here is run on a backend until the sender types its vectors. + */ + @Override + public boolean supportsFeSideResult() { + return false; + } + + /** + * The short circuit produces no Arrow result at either end. PointQueryExecutor is not a + * Coordinator, and Coordinator/NereidsCoordinator are the only places that register a + * FlightSqlEndpointsLocation, so GetFlightInfo found none and failed the query with + * "no FlightSqlEndpointsLocations"; the backend side cannot be pointed at either, since the + * lookup rpc serializes with VMysqlResultWriter into PTabletKeyLookupResponse.row_batch and + * never creates the ArrowFlightResultBlockBuffer that fetch_arrow_flight_schema looks up. + * Arrow Flight SQL stays on the normal execution path. See #67368. + */ + @Override + public boolean supportsShortCircuitPointQuery() { + return false; + } + + /** + * A Flight session does not retry a failed query under a new query id within + * {@code StmtExecutor.handleQueryWithRetry}: the client is not told which of the attempts + * its endpoints belong to. (The replan retry of {@code StmtExecutor.queryRetry} is not asked; + * it starts every attempt through {@link #beforeAttempt}.) + */ + @Override + public boolean canRetryQuery(ConnectContext ctx) { + return false; + } + + /** A statement's result is on this frontend until {@link #beforeQuery} says otherwise. */ + @Override + public void beforeStatement(ConnectContext ctx) { + returnResultFromLocal = true; + endpointsBeforeStatement = endpointsLocations.size(); + } + + /** + * An attempt starts where the statement did: its result is on this frontend, and it has + * registered no endpoint yet. The attempt that failed before it may have moved the result to + * the backends ({@link #beforeQuery}) and registered where; nothing will be pulled from + * there, and a stale "on the backends" state would keep the statement's cleanup (its query + * registration, its connector statement scope) waiting for a DoGet that never comes. Only + * what that attempt registered is withdrawn: what an earlier statement of the request + * registered is left as it was. + */ + @Override + public void beforeAttempt(ConnectContext ctx) { + returnResultFromLocal = true; + if (endpointsLocations.size() > endpointsBeforeStatement) { + endpointsLocations.subList(endpointsBeforeStatement, endpointsLocations.size()).clear(); + } + } + + /** + * The query's result stays on the backends for the client to pull with DoGet; the + * coordinator registers where ({@link #addEndpointsLocation}) instead of fetching the rows. + */ + @Override + public void beforeQuery(ConnectContext ctx) { + returnResultFromLocal = false; + } + + @Override + public boolean returnsResultFromLocal(ConnectContext ctx) { + return returnResultFromLocal; + } + + /** + * The master's response is consumed here as a status and, for a SHOW, a result set (see + * {@link #carryForwardedOutcome}); it is never replayed to the client as packets, so the + * master needs to know nothing about the client. + */ + @Override + public void fillForwardRequest(ConnectContext ctx, TMasterOpRequest request) { + } + @Override public ConnectPoolMgr connectPool(ConnectScheduler scheduler) { return scheduler.getFlightSqlConnectPoolMgr(); @@ -232,16 +335,19 @@ public List getEndpointsLocations() { return endpointsLocations; } - public void clearEndpointsLocations() { + /** + * Starts a request of the session: whatever the previous request left behind is dropped. + * Its query's coordinator, if its close was deferred, is finalized now -- the previous DoGet + * is done by the time the next request arrives (#62259); the result it may have cached and + * never pulled with DoGet is released; its endpoints are forgotten; and the new request's + * result is on this frontend until a query is run for it. + */ + public void beginRequest() { + closeDeferredExecutors(); + channel.reset(); endpointsLocations.clear(); - } - - public void setReturnResultFromLocal(boolean returnResultFromLocal) { - this.returnResultFromLocal = returnResultFromLocal; - } - - public boolean isReturnResultFromLocal() { - return returnResultFromLocal; + endpointsBeforeStatement = 0; + returnResultFromLocal = true; } public void addDeferredExecutor(StmtExecutor executor) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightResultSender.java b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightResultSender.java index f7f31a23119c29..2766a2e31409ee 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightResultSender.java +++ b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightResultSender.java @@ -49,10 +49,6 @@ public class FlightResultSender implements ResultSender { @Override public void sendResultSet(ResultSet resultSet, List fieldInfos, boolean binaryRows) { adapter.getChannel().addResult(DebugUtil.printId(ctx.queryId()), adapter.getRunningQuery(), resultSet); - // The statement's result is on this frontend, whatever the query path decided earlier: an - // EXPLAIN goes through the query path, which marks the result as coming from the backend - // before it knows the statement will not run there. - adapter.setReturnResultFromLocal(true); } @Override @@ -70,6 +66,6 @@ public void sendRow(ByteBuffer row) { @Override public void reset() { // Results are cached per query id and the cache is cleared when the next request of the - // session starts (DorisFlightSqlProducer.executeQueryStatement); nothing is pending here. + // session starts (FlightProtocolAdapter.beginRequest); nothing is pending here. } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/DummyMysqlChannel.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/DummyMysqlChannel.java index 4d738cb639d6a2..bd020e6ecf18ac 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/mysql/DummyMysqlChannel.java +++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/DummyMysqlChannel.java @@ -28,10 +28,6 @@ */ public class DummyMysqlChannel extends MysqlChannel { - public void setSequenceId(int sequenceId) { - this.sequenceId = sequenceId; - } - @Override public String getRemoteIp() { return ""; diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/MysqlChannel.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/MysqlChannel.java index ab321557346797..58f0e93791b179 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/mysql/MysqlChannel.java +++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/MysqlChannel.java @@ -57,6 +57,11 @@ public class MysqlChannel implements BytesChannel { protected static final int SSL_PACKET_HEADER_LEN = 5; // next sequence id to receive or send protected int sequenceId; + // The sequence id the client expects next: sequenceId as of the last packet that reached the + // wire (the command packet received, or the last packet flushed). sequenceId runs ahead of + // it by the packets still in the send buffer; reset() rewinds to it when it drops them, so + // the packets written after a reset are numbered the way the client expects them. + protected int wireSequenceId; // channel connected with client private StreamConnection conn; // used to receive/send header, avoiding new this many time. @@ -145,6 +150,7 @@ public void initSslBuffer() { public void setSequenceId(int sequenceId) { this.sequenceId = sequenceId; + this.wireSequenceId = sequenceId; } public String getRemoteIp() { @@ -429,6 +435,7 @@ public ByteBuffer fetchOnePacket() throws IOException { } if (!isSslHandshaking) { accSequenceId(); + wireSequenceId = sequenceId; } if (packetLen != MAX_PHYSICAL_PACKET_LENGTH) { result.flip(); @@ -504,6 +511,7 @@ public void flush() throws IOException { } finally { sendBuffer.clear(); } + wireSequenceId = sequenceId; isSend = true; } @@ -556,6 +564,7 @@ public void sendOnePacket(ByteBuffer packet) throws IOException { writeHeader(bufLen, isSslMode); writeBuffer(packet); accSequenceId(); + markWireIfSent(); } } if (isSslHandshaking) { @@ -566,6 +575,15 @@ public void sendOnePacket(ByteBuffer packet) throws IOException { packet.limit(oldLimit); writeBuffer(packet); accSequenceId(); + markWireIfSent(); + } + } + + // A packet too large for the send buffer goes to the wire directly (writeBuffer), after what + // the buffer held; when nothing is left in the buffer the client has seen everything written. + private void markWireIfSent() { + if (sendBuffer == null || sendBuffer.position() == 0) { + wireSequenceId = sequenceId; } } @@ -574,12 +592,15 @@ public void sendAndFlush(ByteBuffer packet) throws IOException { flush(); } - // Call this function before send query before + // Drops what was written since the last flush: called before a statement (and before a query is + // attempted again) so that only its own packets reach the client. Those packets are numbered + // from where the client left off, not from where the dropped ones would have ended. public void reset() { isSend = false; if (null != sendBuffer) { sendBuffer.clear(); } + sequenceId = wireSequenceId; } public boolean isSend() { diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/ProxyMysqlChannel.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/ProxyMysqlChannel.java index ed491df4b87e8b..ac31b2ac31e4a0 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/mysql/ProxyMysqlChannel.java +++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/ProxyMysqlChannel.java @@ -34,6 +34,18 @@ public void sendOnePacket(ByteBuffer packet) { proxyResultBuffer.add(packet); } + /** + * Nothing is ever flushed to a client here, so a reset -- before a statement, and before a + * failed query is attempted again -- drops everything the statement wrote so far: the packets + * of a failed attempt must not travel to the client next to those of the attempt that + * succeeded. + */ + @Override + public void reset() { + super.reset(); + proxyResultBuffer.clear(); + } + public List getProxyResultBufferList() { return proxyResultBuffer; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapter.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapter.java index 46f5f35101c96e..d7f997ffc3ee5e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapter.java @@ -20,6 +20,7 @@ import org.apache.doris.common.util.DebugUtil; import org.apache.doris.mysql.MysqlCapability; import org.apache.doris.mysql.MysqlChannel; +import org.apache.doris.mysql.MysqlCommand; import org.apache.doris.mysql.MysqlCursorFetchCompatibility; import org.apache.doris.mysql.MysqlHandshakePacket; import org.apache.doris.mysql.MysqlPacket; @@ -27,6 +28,7 @@ import org.apache.doris.mysql.MysqlSerializer; import org.apache.doris.mysql.MysqlServerStatusFlag; import org.apache.doris.mysql.MysqlSslContext; +import org.apache.doris.mysql.ProxyMysqlChannel; import org.apache.doris.nereids.glue.LogicalPlanAdapter; import org.apache.doris.nereids.stats.StatsErrorEstimator; import org.apache.doris.qe.ConnectContext; @@ -38,6 +40,7 @@ import org.apache.doris.qe.ShowResultSet; import org.apache.doris.qe.StmtExecutor; import org.apache.doris.qe.protocol.ProtocolAdapter; +import org.apache.doris.thrift.TMasterOpRequest; import org.apache.doris.thrift.TResultSinkType; import org.apache.logging.log4j.LogManager; @@ -45,6 +48,7 @@ import java.io.IOException; import java.nio.ByteBuffer; +import java.util.List; /** * The MySQL protocol side of a connection: the channel to the client, the capabilities negotiated @@ -111,6 +115,106 @@ public boolean supportsSqlCacheReplay() { return true; } + @Override + public boolean canReplayForwardedQueryResult() { + return true; + } + + @Override + public boolean supportsFeSideResult() { + return true; + } + + @Override + public boolean supportsShortCircuitPointQuery() { + return true; + } + + /** + * The packets of a failed attempt are dropped when the next attempt resets the channel, so a + * retry is invisible to the client as long as none of them was flushed to the socket yet. + */ + @Override + public boolean canRetryQuery(ConnectContext ctx) { + return !channel.isSend(); + } + + /** + * Clears the send flag and whatever the previous statement of the request left in the send + * buffer. Between the statements of a request run for a client that did not negotiate + * CLIENT_MULTI_STATEMENTS nothing is flushed (see {@link #finishStatement}), so what a + * statement wrote is still in the buffer when the next one starts. + */ + @Override + public void beforeStatement(ConnectContext ctx) { + channel.reset(); + } + + @Override + public void beforeAttempt(ConnectContext ctx) { + // The result of every attempt is relayed by this frontend; a failed attempt registers + // nothing for the client that the next one would have to drop. + } + + @Override + public void beforeQuery(ConnectContext ctx) { + // The rows are relayed through the channel as the coordinator fetches them. + } + + @Override + public boolean returnsResultFromLocal(ConnectContext ctx) { + return true; + } + + /** + * The master encodes the response of a forwarded statement for this connection's client, so + * it needs the client's negotiated capabilities, and for a forwarded COM_STMT_EXECUTE the + * execute packet and whether it asked for a cursor. {@link #restoreFromForwardRequest} is the + * master's side. + */ + @Override + public void fillForwardRequest(ConnectContext ctx, TMasterOpRequest request) { + if (ctx.getCommand() == MysqlCommand.COM_STMT_EXECUTE) { + if (prepareExecuteBuffer != null) { + request.setPrepareExecuteBuffer(prepareExecuteBuffer); + } + request.setCursorFetchRequested(cursorFetchRequested); + } + request.setClientDeprecatedEOF(channel.clientDeprecatedEOF()); + request.setMysqlCapability(capability.getFlags()); + } + + /** + * On the master, gives the proxy context of a forwarded statement the capabilities of the + * client it is answering, as {@link #fillForwardRequest} put them in the request. A request + * from an old frontend carries neither the capability flags nor the cursor flag; it gets the + * default capabilities without CLIENT_DEPRECATE_EOF, plus that flag when set separately, and + * its ordinary prepared statements are not rejected. + */ + public void restoreFromForwardRequest(ConnectContext ctx, TMasterOpRequest request) { + int flags = request.isSetMysqlCapability() ? request.getMysqlCapability() + : MysqlCapability.DEFAULT_CAPABILITY.getFlags() + & ~MysqlCapability.Flag.CLIENT_DEPRECATE_EOF.getFlagBit(); + if (request.isSetClientDeprecatedEOF() && request.isClientDeprecatedEOF()) { + flags |= MysqlCapability.Flag.CLIENT_DEPRECATE_EOF.getFlagBit(); + } + MysqlCapability restored = new MysqlCapability(flags); + setCapability(restored); + channel.getSerializer().setCapability(restored); + if (restored.isDeprecatedEOF()) { + channel.setClientDeprecatedEOF(); + } + cursorFetchRequested = request.isSetCursorFetchRequested() && request.isCursorFetchRequested(); + } + + /** + * On the master, the packets the forwarded statement produced, collected by the proxy + * context's channel to be handed back to the frontend the client is connected to. + */ + public List proxyResultPackets() { + return ((ProxyMysqlChannel) channel).getProxyResultBufferList(); + } + @Override public ConnectPoolMgr connectPool(ConnectScheduler scheduler) { return scheduler.getConnectPoolMgr(); @@ -120,9 +224,10 @@ public ConnectPoolMgr connectPool(ConnectScheduler scheduler) { * Between the statements of a multi-statement request the intermediate response carries * SERVER_MORE_RESULTS_EXISTS, and is sent right away if the client negotiated * CLIENT_MULTI_STATEMENTS. Here Doris differs from MySQL: a client that did not negotiate it - * gets the request run as several statements anyway, but only the last result is delivered - * (the next query resets the channel, see {@link MysqlResultSender#reset}). The response of the - * last statement is the response of the command, sent by {@link #finishCommand}. + * gets the request run as several statements anyway, but only the last statement's outcome + * is delivered (the next statement resets the channel, see {@link #beforeStatement}). The + * response of the last statement is the response of the command, sent by + * {@link #finishCommand}. */ @Override public boolean finishStatement(ConnectContext ctx, StmtExecutor executor, int stmtIndex, int stmtCount) diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlResultSender.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlResultSender.java index 9824cd07485a7e..7ff8cfd257f3da 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlResultSender.java +++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlResultSender.java @@ -138,7 +138,7 @@ public void sendRow(ByteBuffer row) throws IOException { channel().sendOnePacket(row); } - /** Clears the send flag and whatever the previous statement left in the send buffer. */ + /** Clears the send flag and whatever a failed attempt of the query left in the send buffer. */ @Override public void reset() { channel().reset(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java index 51bdc44b66b9c6..140dc8a603d1c3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/LogicalResultSinkToShortCircuitPointQuery.java @@ -31,7 +31,6 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.qe.ConnectContext; -import org.apache.doris.qe.ConnectContext.ConnectType; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; @@ -69,16 +68,12 @@ boolean scanMatchShortCircuitCondition(LogicalOlapScan olapScan) { if (!connectContext.getSessionVariable().isEnableShortCircuitQuery()) { return false; } - // The short circuit produces no Arrow result at either end. PointQueryExecutor is not a - // Coordinator, and Coordinator/NereidsCoordinator are the only places that register a - // FlightSqlEndpointsLocation, so GetFlightInfo found none and failed the query with - // "no FlightSqlEndpointsLocations"; the BE side cannot be pointed at either, since the lookup rpc - // serializes with VMysqlResultWriter into PTabletKeyLookupResponse.row_batch and never creates the - // ArrowFlightResultBlockBuffer that fetch_arrow_flight_schema looks up. Keep Arrow Flight SQL on - // the normal execution path. This has to be decided here at plan time rather than when picking the - // executor: OlapScanNode.computeTabletInfo and several rewrite and property rules read - // StatementContext.isShortCircuitQuery() while building the plan. See #67368. - if (connectContext.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) { + // A protocol whose client pulls the result from the backend has no result to pull for a + // short circuit (see FlightProtocolAdapter.supportsShortCircuitPointQuery). This has to be + // decided here at plan time rather than when picking the executor: OlapScanNode.computeTabletInfo + // and several rewrite and property rules read StatementContext.isShortCircuitQuery() while + // building the plan. See #67368. + if (!connectContext.getProtocolAdapter().supportsShortCircuitPointQuery()) { return false; } // Lazy point-query pruning does not preserve explicit PARTITION/TABLET restrictions. diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/StartTransactionCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/StartTransactionCommand.java index c4458934f942b8..1286b81a8eb746 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/StartTransactionCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/StartTransactionCommand.java @@ -32,9 +32,6 @@ public StartTransactionCommand() { @Override public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { - if (ctx.getConnectType() == ConnectContext.ConnectType.MYSQL) { - ctx.getMysqlChannel().reset(); - } // do nothing ctx.getState().setOk(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionBeginCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionBeginCommand.java index a76bd5cd5ed5f7..d36e9a98085cf6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionBeginCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionBeginCommand.java @@ -64,10 +64,6 @@ private void validate(ConnectContext ctx) { } private void handleTransactionBegin(ConnectContext ctx) { - if (ctx.getConnectType() == ConnectContext.ConnectType.MYSQL) { - // Every time set no send flag and clean all data in buffer - ctx.getMysqlChannel().reset(); - } ctx.getState().setOk(0, 0, ""); if (ctx.isTxnModel()) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionCommitCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionCommitCommand.java index 70580d1ee66e39..712c230c2a5fa1 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionCommitCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionCommitCommand.java @@ -45,10 +45,6 @@ public void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception { } private void handleTransactionCommit(ConnectContext ctx) throws AnalysisException { - if (ctx.getConnectType() == ConnectContext.ConnectType.MYSQL) { - // Every time set no send flag and clean all data in buffer - ctx.getMysqlChannel().reset(); - } ctx.getState().setOk(0, 0, ""); if (ctx.getTxnEntry() != null && ctx.getTxnEntry().getRowsInTransaction() == 0 diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionRollbackCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionRollbackCommand.java index 173ed777c7cce1..15b58a4ced9555 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionRollbackCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/TransactionRollbackCommand.java @@ -45,10 +45,6 @@ public void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception { } private void handleTransactionRollback(ConnectContext ctx) throws AnalysisException { - if (ctx.getConnectType() == ConnectContext.ConnectType.MYSQL) { - // Every time set no send flag and clean all data in buffer - ctx.getMysqlChannel().reset(); - } ctx.getState().setOk(0, 0, ""); if (ctx.getTxnEntry() != null && ctx.getTxnEntry().getRowsInTransaction() == 0 diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/BatchInsertIntoTableCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/BatchInsertIntoTableCommand.java index b99eb39e296497..55ddf9fb4ece81 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/BatchInsertIntoTableCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/BatchInsertIntoTableCommand.java @@ -48,7 +48,6 @@ import org.apache.doris.nereids.trees.plans.physical.PhysicalUnion; import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; import org.apache.doris.qe.ConnectContext; -import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.qe.StmtExecutor; import com.google.common.base.Preconditions; @@ -137,9 +136,6 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { statementContext, supportFastInsertIntoValues, true); planner.plan(logicalPlanAdapter, ctx.getSessionVariable().toThrift()); executor.checkBlockRules(); - if (ctx.getConnectType() == ConnectType.MYSQL && ctx.getMysqlChannel() != null) { - ctx.getMysqlChannel().reset(); - } Optional> plan = planner.getPhysicalPlan() .>collect(PhysicalOlapTableSink.class::isInstance).stream().findAny(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTVFCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTVFCommand.java index 6c8826c6dae76b..30cd5da46e6218 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTVFCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTVFCommand.java @@ -40,7 +40,6 @@ import org.apache.doris.nereids.trees.plans.physical.PhysicalTVFTableSink; import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; import org.apache.doris.qe.ConnectContext; -import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.qe.Coordinator; import org.apache.doris.qe.QeProcessorImpl; import org.apache.doris.qe.QeProcessorImpl.QueryInfo; @@ -113,10 +112,6 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { } } - if (ctx.getConnectType() == ConnectType.MYSQL && ctx.getMysqlChannel() != null) { - ctx.getMysqlChannel().reset(); - } - // 3. Create coordinator Coordinator coordinator = EnvFactory.getInstance().createCoordinator( ctx, planner, ctx.getStatsErrorEstimator()); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTableCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTableCommand.java index ff99b889f16a90..0882a19a7a0624 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTableCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTableCommand.java @@ -82,7 +82,6 @@ import org.apache.doris.planner.PlanFragment; import org.apache.doris.planner.PlanNode; import org.apache.doris.qe.ConnectContext; -import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.qe.Coordinator; import org.apache.doris.qe.StmtExecutor; import org.apache.doris.system.Backend; @@ -516,9 +515,6 @@ ExecutorFactory selectInsertExecutorFactory( try { stmtExecutor.setPlanner(planner); stmtExecutor.checkBlockRules(); - if (ctx.getConnectType() == ConnectType.MYSQL && ctx.getMysqlChannel() != null) { - ctx.getMysqlChannel().reset(); - } Optional> plan = (planner.getPhysicalPlan() .>collect(PhysicalSink.class::isInstance)).stream() .findAny(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteTableCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteTableCommand.java index bf752f214c9f5f..5d67020025d797 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteTableCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteTableCommand.java @@ -62,7 +62,6 @@ import org.apache.doris.nereids.trees.plans.physical.PhysicalTableSink; import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; import org.apache.doris.qe.ConnectContext; -import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.qe.QueryState.MysqlStateType; import org.apache.doris.qe.StmtExecutor; import org.apache.doris.thrift.TPartialUpdateNewRowPolicy; @@ -180,9 +179,6 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { Plan analyzedPlan = planner.getAnalyzedPlan(); lineagePlan = Optional.ofNullable(analyzedPlan); executor.checkBlockRules(); - if (ctx.getConnectType() == ConnectType.MYSQL && ctx.getMysqlChannel() != null) { - ctx.getMysqlChannel().reset(); - } Optional> plan = (planner.getPhysicalPlan() .>collect(node -> node instanceof PhysicalTableSink)).stream().findAny(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/RewriteTableCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/RewriteTableCommand.java index f6af9478c894f0..68ff5feadc520f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/RewriteTableCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/RewriteTableCommand.java @@ -175,9 +175,6 @@ private ExecutorFactory selectInsertExecutorFactory(NereidsPlanner planner, Conn try { stmtExecutor.setPlanner(planner); stmtExecutor.checkBlockRules(); - if (ctx.getConnectType() == ConnectContext.ConnectType.MYSQL && ctx.getMysqlChannel() != null) { - ctx.getMysqlChannel().reset(); - } Optional> plan = (planner.getPhysicalPlan() .>collect(PhysicalSink.class::isInstance)).stream().findAny(); Preconditions.checkArgument(plan.isPresent(), "rewrite command must contain target table"); 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 92a6a17bdc42e5..829f28714f362b 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 @@ -903,19 +903,12 @@ public List getFlightSqlEndpointsLocations() { return FlightProtocolAdapter.of(this).getEndpointsLocations(); } - public void clearFlightSqlEndpointsLocations() { - FlightProtocolAdapter.of(this).clearEndpointsLocations(); - } - - public void setReturnResultFromLocal(boolean returnResultFromLocal) { - FlightProtocolAdapter.of(this).setReturnResultFromLocal(returnResultFromLocal); - } - - // A MySQL connection always sends its result from this frontend; only an Arrow Flight SQL - // session may leave a query's result on the backend for the client to pull. + /** + * Whether the result of the statement being executed comes from this frontend, or is left + * on the backends for the client to pull; see {@link ProtocolAdapter#returnsResultFromLocal}. + */ public boolean isReturnResultFromLocal() { - return !(protocolAdapter instanceof FlightProtocolAdapter) - || ((FlightProtocolAdapter) protocolAdapter).isReturnResultFromLocal(); + return protocolAdapter.returnsResultFromLocal(this); } // The bearer token of an Arrow Flight SQL session, null for any other connection. @@ -1509,7 +1502,7 @@ public String getQueryIdentifier() { } public boolean supportHandleByFe() { - return !getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL) && getCommand() != MysqlCommand.COM_STMT_EXECUTE; + return protocolAdapter.supportsFeSideResult() && getCommand() != MysqlCommand.COM_STMT_EXECUTE; } public void setCloudCluster(String cluster) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectProcessor.java index 389825b74f34e9..ad47363d882576 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectProcessor.java @@ -52,7 +52,6 @@ import org.apache.doris.datasource.DelegatedCredential; import org.apache.doris.datasource.SessionContext; import org.apache.doris.metric.MetricRepo; -import org.apache.doris.mysql.MysqlCapability; import org.apache.doris.mysql.MysqlCommand; import org.apache.doris.mysql.MysqlServerStatusFlag; import org.apache.doris.mysql.protocol.MysqlProtocolAdapter; @@ -322,6 +321,7 @@ public void executeQuery(String originStmt) throws Exception { if (i > 0) { ctx.resetReturnRows(); } + ctx.getProtocolAdapter().beforeStatement(ctx); // Re-resolve per statement: an earlier statement in the same multi-stmt // request (e.g. SET workload_group=...) may have changed the effective // workload group, and later statements that fail before Coordinator.exec @@ -578,7 +578,7 @@ public TMasterOpResult proxyExecute(TMasterOpRequest request) throws TException ctx.setThreadLocalInfo(); StmtExecutor executor = null; try { - restoreForwardedMysqlContext(ctx, request); + MysqlProtocolAdapter.of(ctx).restoreFromForwardRequest(ctx, request); // 0 for compatibility. int idx = request.isSetStmtIdx() ? request.getStmtIdx() : 0; executor = new StmtExecutor(ctx, new OriginStatement(request.getSql(), idx), true); @@ -694,31 +694,13 @@ public TMasterOpResult proxyExecute(TMasterOpRequest request) throws TException } if (executor.getProxyShowResultSet() != null) { result.setResultSet(executor.getProxyShowResultSet().tothrift()); - } else if (!executor.getProxyQueryResultBufList().isEmpty()) { - result.setQueryResultBufList(executor.getProxyQueryResultBufList()); + } else if (!mysqlAdapter.proxyResultPackets().isEmpty()) { + result.setQueryResultBufList(mysqlAdapter.proxyResultPackets()); } } return result; } - static void restoreForwardedMysqlContext(ConnectContext context, TMasterOpRequest request) { - int flags = request.isSetMysqlCapability() ? request.getMysqlCapability() - : MysqlCapability.DEFAULT_CAPABILITY.getFlags() - & ~MysqlCapability.Flag.CLIENT_DEPRECATE_EOF.getFlagBit(); - if (request.isSetClientDeprecatedEOF() && request.isClientDeprecatedEOF()) { - flags |= MysqlCapability.Flag.CLIENT_DEPRECATE_EOF.getFlagBit(); - } - MysqlCapability capability = new MysqlCapability(flags); - context.setCapability(capability); - context.getMysqlChannel().getSerializer().setCapability(capability); - if (capability.isDeprecatedEOF()) { - context.getMysqlChannel().setClientDeprecatedEOF(); - } - // Old followers do not carry the cursor flag. Keep their existing behavior; they must - // be upgraded to preserve cursor intent. Do not reject their ordinary prepared statements. - context.setCursorFetchRequested(request.isSetCursorFetchRequested() && request.isCursorFetchRequested()); - } - static void restoreForwardedSessionContext(ConnectContext context, TMasterOpRequest request) { if (!request.isSetDelegatedCredentialToken()) { return; diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java index 0d3dec78561108..b3005119418039 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java @@ -78,7 +78,6 @@ import org.apache.doris.proto.InternalService.PExecPlanFragmentStartRequest; import org.apache.doris.proto.Types; import org.apache.doris.proto.Types.PUniqueId; -import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.qe.QueryStatisticsItem.FragmentInstanceInfo; import org.apache.doris.resource.BackendSelection; import org.apache.doris.resource.BackendSelectionManager; @@ -877,7 +876,7 @@ protected void execInternal() throws Exception { toBrpcHost(param.host), this.timeoutDeadline, context.getSessionVariable().getMaxMsgSizeOfResultReceiver(), enableParallelResultSink)); } else { - Preconditions.checkState(context.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)); + // The client pulls the result from the backend (Arrow Flight SQL); register where. TUniqueId finstId; if (enableParallelResultSink) { finstId = queryId; diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/FEOpExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/FEOpExecutor.java index 11503c1b895b40..39828f852d7c21 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/FEOpExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/FEOpExecutor.java @@ -25,12 +25,10 @@ import org.apache.doris.common.Config; import org.apache.doris.common.ErrorCode; import org.apache.doris.datasource.DelegatedCredential; -import org.apache.doris.mysql.MysqlCommand; import org.apache.doris.mysql.MysqlProto; import org.apache.doris.mysql.MysqlResultSetEndPacket; import org.apache.doris.mysql.MysqlSerializer; import org.apache.doris.mysql.protocol.MysqlProtocolAdapter; -import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.thrift.FrontendService; import org.apache.doris.thrift.TExpr; import org.apache.doris.thrift.TExprNode; @@ -216,13 +214,6 @@ protected TMasterOpRequest buildStmtForwardParams() throws AnalysisException { params.setTxnLoadInfo(ctx.getTxnEntry().getTxnLoadInfoInObserver()); } - if (ctx.getCommand() == MysqlCommand.COM_STMT_EXECUTE) { - if (null != ctx.getPrepareExecuteBuffer()) { - params.setPrepareExecuteBuffer(ctx.getPrepareExecuteBuffer()); - } - params.setCursorFetchRequested(ctx.isCursorFetchRequested()); - } - ctx.getSessionContext().getDelegatedCredential().ifPresent((DelegatedCredential credential) -> { params.setDelegatedCredentialSessionId(ctx.getSessionContext().getSessionId()); params.setDelegatedCredentialType(credential.getType().name()); @@ -230,15 +221,10 @@ protected TMasterOpRequest buildStmtForwardParams() throws AnalysisException { credential.getExpiresAtMillis().ifPresent(params::setDelegatedCredentialExpiresAtMillis); }); - // Propagate the client's CLIENT_DEPRECATE_EOF capability so the master FE - // generates packets matching the original client's protocol expectations. - // Only a MySQL connection negotiates this capability and owns a MysqlChannel; - // an Arrow Flight SQL session has none, and leaving the field unset keeps the - // master on its default packet layout. - if (ctx.getConnectType() == ConnectType.MYSQL) { - params.setClientDeprecatedEOF(ctx.getMysqlChannel().clientDeprecatedEOF()); - params.setMysqlCapability(ctx.getCapability().getFlags()); - } + // What the master needs to know about the client to produce the response it expects: the + // negotiated capabilities and the COM_STMT_EXECUTE packet for a MySQL connection, nothing + // for a protocol that consumes the master's status and rows rather than its packets. + ctx.getProtocolAdapter().fillForwardRequest(ctx, params); return params; } @@ -285,13 +271,14 @@ public boolean hasQueryResultPackets() { // result at the follower, which still has the original execute flag and client capability. // DML/DDL OK and ERR packets are retained verbatim, including warnings and load info. public void prepareQueryResultForClient() { - if (!ctx.getMysqlChannel().clientDeprecatedEOF() || isClientDeprecatedEofApplied() + MysqlProtocolAdapter mysqlAdapter = MysqlProtocolAdapter.of(ctx); + if (!mysqlAdapter.getChannel().clientDeprecatedEOF() || isClientDeprecatedEofApplied() || !hasQueryResultPackets()) { return; } List packets = new ArrayList<>(result.getQueryResultBufList()); int metadataEnd = Math.toIntExact(MysqlProto.readVInt(packets.get(0).duplicate())) + 1; - boolean needsCursorTerminator = MysqlProtocolAdapter.of(ctx).clientConsumesCursorMetadataTerminator(ctx); + boolean needsCursorTerminator = mysqlAdapter.clientConsumesCursorMetadataTerminator(ctx); // An execution error may occur after only part of the metadata has been buffered. if (metadataEnd > packets.size()) { Preconditions.checkState(isErrorPacket(result.packet)); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java index 3cbecace7a2b66..cf9504d852ff36 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/MysqlConnectProcessor.java @@ -129,6 +129,7 @@ protected void handleExecute(PrepareCommand prepareCommand, long stmtId, Prepare // null bitmap String stmtStr = ""; try { + ctx.getProtocolAdapter().beforeStatement(ctx); StatementContext statementContext = prepCtx.getStatementContext(); if (!ctx.isProxy()) { // An empty buffer still identifies a zero-parameter COM_STMT_EXECUTE when forwarding. diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java index af96302568a4de..8c53bca65f7ce9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java @@ -42,7 +42,6 @@ import org.apache.doris.planner.ResultSink; import org.apache.doris.planner.ScanNode; import org.apache.doris.planner.SchemaScanNode; -import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.qe.QueryStatisticsItem.FragmentInstanceInfo; import org.apache.doris.qe.runtime.LoadProcessor; import org.apache.doris.qe.runtime.MultiFragmentsPipelineTask; @@ -507,8 +506,8 @@ private void setForArrowFlight(CoordinatorContext coordinatorContext, PipelineDi ConnectContext connectContext = coordinatorContext.connectContext; DataSink dataSink = coordinatorContext.dataSink; if (dataSink instanceof ResultSink || dataSink instanceof ResultFileSink) { + // The client pulls the result from the backend (Arrow Flight SQL); register where. if (connectContext != null && !connectContext.isReturnResultFromLocal()) { - Preconditions.checkState(connectContext.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)); for (AssignedJob instance : topPlan.getInstanceJobs()) { BackendWorker worker = (BackendWorker) instance.getAssignedWorker(); Backend backend = worker.getBackend(); 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 7f71454ec166c5..9dd0c5f637e2b6 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 @@ -70,7 +70,6 @@ import org.apache.doris.metric.MetricRepo; import org.apache.doris.mysql.FieldInfo; import org.apache.doris.mysql.MysqlCommand; -import org.apache.doris.mysql.ProxyMysqlChannel; import org.apache.doris.mysql.protocol.MysqlProtocolAdapter; import org.apache.doris.nereids.NereidsPlanner; import org.apache.doris.nereids.PlanProcess; @@ -120,7 +119,6 @@ import org.apache.doris.proto.InternalService.POutfileWriteSuccessRequest; import org.apache.doris.proto.InternalService.POutfileWriteSuccessResult; import org.apache.doris.qe.CommonResultSet.CommonResultSetMetaData; -import org.apache.doris.qe.ConnectContext.ConnectType; import org.apache.doris.qe.QeProcessorImpl.QueryInfo; import org.apache.doris.qe.QueryState.MysqlStateType; import org.apache.doris.qe.cache.Cache; @@ -696,9 +694,10 @@ public void execute(TUniqueId queryId) throws Exception { SessionVariable sessionVariable = context.getSessionVariable(); context.setEffectiveCloudCluster(null); externalDmlAuditCoordinator = null; - if (context.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) { - context.setReturnResultFromLocal(true); - } + // Every attempt starts as the statement did. queryRetry() runs this again after a replan + // error without going back through the processor's beforeStatement, and the failed attempt + // may have moved the result to the backends (beforeQuery) and registered where. + context.getProtocolAdapter().beforeAttempt(context); try { try { @@ -902,12 +901,12 @@ private void executeByNereids(TUniqueId queryId) throws Exception { if (context.getCommand() == MysqlCommand.COM_STMT_PREPARE) { throw new UserException("Forward master command is not supported for prepare statement"); } - if (context.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) { + if (!context.getProtocolAdapter().canReplayForwardedQueryResult()) { // The master returns a query result as MySQL wire packets in - // TMasterOpResult.queryResultBufList, which only ConnectProcessor.finalizeCommand() - // can replay and which cannot be converted to Arrow batches. Refuse here, before - // the RPC, rather than let the master build a result set this FE would discard - // and answer the client with a synthesized empty success. + // TMasterOpResult.queryResultBufList, which only a MySQL connection can replay to + // its client (MysqlProtocolAdapter.finishCommand). Refuse here, before the RPC, + // rather than let the master build a result set this FE would discard and answer + // the client with a synthesized empty success. throw new UserException("Forwarding a query to the master FE is not supported on an" + " Arrow Flight SQL connection. Connect to the master FE to run this query."); } @@ -1150,9 +1149,6 @@ private void handleQueryWithRetry(TUniqueId queryId) throws Exception { } } } - if (context.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) { - context.setReturnResultFromLocal(false); - } handleQueryStmt(); LOG.info("Query {} finished", DebugUtil.printId(context.queryId)); break; @@ -1202,8 +1198,7 @@ private void handleQueryWithRetry(TUniqueId queryId) throws Exception { } } } - if (i != retryTime - 1 && isNeedRetry - && context.getConnectType().equals(ConnectType.MYSQL) && !context.getMysqlChannel().isSend()) { + if (i != retryTime - 1 && isNeedRetry && context.getProtocolAdapter().canRetryQuery(context)) { LOG.warn("retry {} times. stmt: {}", (i + 1), parsedStmt.getOrigStmt().originStmt); } else { throw e; @@ -1448,7 +1443,9 @@ private void handleQueryStmt() throws Exception { } ResultSender sender = context.getResultSender(); - // Every time set no send flag and clean all data in buffer + // Each attempt of the query starts from a clean sender: a failed attempt may have left + // packets behind that never reached the client (or it would not be retried, see + // ProtocolAdapter.canRetryQuery). sender.reset(); Queriable queryStmt = (Queriable) parsedStmt; @@ -1531,6 +1528,11 @@ public void executeAndSendResult(boolean isOutfileQuery, boolean isSendFields, // Query OK, 10 rows affected (0.01 sec) // // 2. If this is a query, send the result expr fields first, and send result data back to client. + // + // Where the result goes is the protocol's decision, made now, before the coordinator is + // built: relayed by this frontend through the sender, or left on the backends for the + // client to pull (context.isReturnResultFromLocal() is false then). + context.getProtocolAdapter().beforeQuery(context); RowBatch batch; CoordInterface coordBase = null; if (statementContext.isShortCircuitQuery()) { @@ -1574,22 +1576,21 @@ public void executeAndSendResult(boolean isOutfileQuery, boolean isSendFields, profile.getSummaryProfile().setQueryScheduleFinishTime(TimeUtils.getStartTimeMs()); updateProfile(false); - if (context.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)) { - Preconditions.checkState(!context.isReturnResultFromLocal()); + if (!context.isReturnResultFromLocal()) { profile.getSummaryProfile().setTempStartTime(); - // The client pulls the results from the BE later (DoGet). Only an external-table - // scan in batch mode still needs the coordinator after this point: the BE fetches - // its splits lazily from the split source the coordinator holds, so closing the - // coordinator here would release that source too early and break DoGet (#62259). - // Such a coordinator is closed later by ConnectContext: on the session's next - // query, on teardown, or by the idle reaper in checkTimeout. The trade-off is that - // its query queue slot and query registration stay held until then. Every other - // query closes its coordinator in the finally block below and releases both right - // away, the BE buffering its results independently of the coordinator (#67503). - // A short-circuit point query is the one case with a different coordBase, and it - // can no longer reach here: it has no Arrow result on either side, so - // LogicalResultSinkToShortCircuitPointQuery keeps Arrow Flight SQL on the normal - // execution path (#67368). + // The client pulls the results from the BE later (Arrow Flight SQL's DoGet). Only an + // external-table scan in batch mode still needs the coordinator after this point: + // the BE fetches its splits lazily from the split source the coordinator holds, so + // closing the coordinator here would release that source too early and break DoGet + // (#62259). Such a coordinator is closed later by ConnectContext: on the session's + // next query, on teardown, or by the idle reaper in checkTimeout. The trade-off is + // that its query queue slot and query registration stay held until then. Every + // other query closes its coordinator in the finally block below and releases both + // right away, the BE buffering its results independently of the coordinator + // (#67503). A short-circuit point query is the one case with a different coordBase, + // and it cannot reach here: it has no Arrow result on either side, so + // LogicalResultSinkToShortCircuitPointQuery keeps a Flight session on the normal + // execution path (ProtocolAdapter.supportsShortCircuitPointQuery, #67368). if (coordBase == coord && coord.hasBatchSplitSource()) { deferForArrowFlight(); } @@ -2280,10 +2281,6 @@ private String getStmtForLoggingBeforeParse(String stmt) { } } - public List getProxyQueryResultBufList() { - return ((ProxyMysqlChannel) context.getMysqlChannel()).getProxyResultBufferList(); - } - public void sendProxyQueryResult() throws IOException { if (masterOpExecutor == null) { return; diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ProtocolAdapter.java b/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ProtocolAdapter.java index 4b8806575d7384..94c418ee5a21b4 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ProtocolAdapter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ProtocolAdapter.java @@ -22,6 +22,7 @@ import org.apache.doris.qe.ConnectPoolMgr; import org.apache.doris.qe.ConnectScheduler; import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.thrift.TMasterOpRequest; import org.apache.doris.thrift.TResultSinkType; import java.io.IOException; @@ -62,6 +63,78 @@ public interface ProtocolAdapter { */ boolean supportsSqlCacheReplay(); + /** + * Whether a query forwarded to the master can be answered with the result the master sends + * back. The master produces it as MySQL wire packets ({@code TMasterOpResult.queryResultBufList}), + * so only a MySQL connection can replay it; a connection of any other protocol refuses to + * forward a query instead of answering it with a synthesized empty success. Forwarded + * statements whose result is a {@code ShowResultSet} or just a status are not affected. + */ + boolean canReplayForwardedQueryResult(); + + /** + * Whether a query the planner can answer without a backend (a literal, a session variable) + * may be answered by this frontend, through {@link ResultSender#sendResultSet}, instead of + * being run on a backend. + */ + boolean supportsFeSideResult(); + + /** + * Whether a point query on a merge-on-write unique table may take the short circuit: + * {@code PointQueryExecutor} looks the row up on the backend over a plain rpc and encodes it + * on this frontend, instead of running a query whose result the backend keeps in a result + * sink of this connection's {@link #resultSinkType}. + */ + boolean supportsShortCircuitPointQuery(); + + /** + * Whether the query being executed, which just failed, may be run again under a new query id + * without the client noticing. Asked after each failed attempt of + * {@code StmtExecutor.handleQueryWithRetry}. + */ + boolean canRetryQuery(ConnectContext ctx); + + /** + * Called by the connect processor before each statement of a request is executed. The + * protocol drops what the previous statement of the same request left behind, so that a + * request delivers only the outcome of its last statement. + */ + void beforeStatement(ConnectContext ctx); + + /** + * Called by the executor at the start of every attempt to execute the statement: the first + * one, and each one {@code StmtExecutor.queryRetry} makes after an attempt failed with an + * error the statement is replanned on. A retried attempt does not go through + * {@link #beforeStatement}, so the protocol drops here what the failed attempt left behind + * and starts the new one as the statement started: its result comes from this frontend until + * {@link #beforeQuery} moves it, and nothing the failed attempt registered for the client is + * delivered. What earlier statements of the request left is not touched. + */ + void beforeAttempt(ConnectContext ctx); + + /** + * Called by the executor when the statement's plan is about to be run on the backends as a + * query, before the coordinator is built. The protocol decides here where the query's result + * goes: relayed by this frontend row by row, or left on the backends for the client to pull; + * see {@link #returnsResultFromLocal}. + */ + void beforeQuery(ConnectContext ctx); + + /** + * Whether the result of the statement being executed comes from this frontend: materialized + * by it, or relayed by it from the backends, and delivered through the {@link ResultSender}. + * It does not when the backends keep the result for the client to pull; such a statement is + * not over when its executor returns, and its query registration and coordinator are released + * only once the client has pulled the result. + */ + boolean returnsResultFromLocal(ConnectContext ctx); + + /** + * Adds to the request that forwards a statement to the master what the master needs to + * know about this connection's client to produce the response the client expects. + */ + void fillForwardRequest(ConnectContext ctx, TMasterOpRequest request); + /** * Called by {@code ConnectProcessor.executeQuery} after the {@code stmtIndex}-th of the * {@code stmtCount} statements of one request has been executed, before it is audited. The diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ResultSender.java b/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ResultSender.java index 36d6f04c59d16c..1e4ee2e14d85a4 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ResultSender.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ResultSender.java @@ -67,8 +67,9 @@ public interface ResultSender { void sendRow(ByteBuffer row) throws IOException; /** - * Forgets whatever the previous statement of the same request left unsent, so that a - * multi-statement request delivers only the last result. Called when a query starts. + * Forgets whatever a failed attempt of the query left unsent, so that the next attempt starts + * from nothing. Called when each attempt of a query starts; what a previous statement of the + * same request left behind is dropped earlier, by {@link ProtocolAdapter#beforeStatement}. */ void reset(); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlProducerTest.java b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlProducerTest.java index 60d146ce41f7b3..f72918e4aa26d9 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlProducerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlProducerTest.java @@ -17,6 +17,7 @@ package org.apache.doris.arrowflight; +import org.apache.doris.arrowflight.protocol.FlightProtocolAdapter; import org.apache.doris.arrowflight.results.FlightSqlChannel; import org.apache.doris.arrowflight.sessions.FlightSessionsManager; import org.apache.doris.common.FeConstants; @@ -171,7 +172,7 @@ public void testGetFlightInfoFinalizesDeferredExecutorWhenSchemaFetchFails() thr // handleQuery plans + submits to BE and defers the coordinator (coordBase == coord), // exactly as executeAndSendResult() does for an Arrow Flight external-table scan. Mockito.doAnswer(invocation -> { - ctx.setReturnResultFromLocal(false); + FlightProtocolAdapter.of(ctx).beforeQuery(ctx); ctx.addFlightSqlDeferredExecutor(deferred); return null; }).when(mock).handleQuery(Mockito.anyString()); diff --git a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java index 9d54b996bf29f0..653b085912761a 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java @@ -17,6 +17,7 @@ package org.apache.doris.arrowflight.protocol; +import org.apache.doris.arrowflight.results.FlightSqlEndpointsLocation; import org.apache.doris.catalog.Column; import org.apache.doris.catalog.ScalarType; import org.apache.doris.common.ErrorCode; @@ -30,6 +31,8 @@ import org.apache.doris.qe.ShowResultSet; import org.apache.doris.qe.ShowResultSetMetaData; import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.thrift.TMasterOpRequest; +import org.apache.doris.thrift.TNetworkAddress; import org.apache.doris.thrift.TResultSinkType; import org.apache.doris.thrift.TUniqueId; @@ -218,6 +221,128 @@ public void testFailedCommandReleasesTheSession() throws Exception { Assertions.assertEquals("ok", callFromAnotherThread(adapter, ctx)); } + @Test + public void testAFlightSessionTakesNoResultProducedForAMysqlClient() { + ConnectContext ctx = flightSession(); + FlightProtocolAdapter adapter = FlightProtocolAdapter.of(ctx); + + // The SQL cache and the master answer with MySQL packets; a FE-side result is untyped + // Utf8; a short circuit has no Arrow result at either end; a retry would leave the failed + // attempt's endpoints behind. + Assertions.assertFalse(adapter.supportsSqlCacheReplay()); + Assertions.assertFalse(adapter.canReplayForwardedQueryResult()); + Assertions.assertFalse(adapter.supportsFeSideResult()); + Assertions.assertFalse(ctx.supportHandleByFe()); + Assertions.assertFalse(adapter.supportsShortCircuitPointQuery()); + Assertions.assertFalse(adapter.canRetryQuery(ctx)); + + // The master needs to know nothing about the client: its response is consumed here. + TMasterOpRequest request = new TMasterOpRequest(); + adapter.fillForwardRequest(ctx, request); + Assertions.assertFalse(request.isSetMysqlCapability()); + Assertions.assertFalse(request.isSetClientDeprecatedEOF()); + Assertions.assertFalse(request.isSetPrepareExecuteBuffer()); + Assertions.assertFalse(request.isSetCursorFetchRequested()); + } + + @Test + public void testWhereTheResultIsFollowsTheStatementLifecycle() throws Exception { + ConnectContext ctx = flightSession(); + FlightProtocolAdapter adapter = FlightProtocolAdapter.of(ctx); + ShowResultSet resultSet = new ShowResultSet( + ShowResultSetMetaData.builder().addColumn(new Column("c", ScalarType.createVarchar(20))).build(), + Lists.>newArrayList(Lists.newArrayList("v"))); + + // A statement's result is on this frontend (a SHOW, a SET) ... + adapter.beforeStatement(ctx); + Assertions.assertTrue(ctx.isReturnResultFromLocal()); + // ... until a query is run for it on the backends: then the client pulls it from where + // the coordinator registered it. + adapter.beforeQuery(ctx); + Assertions.assertFalse(ctx.isReturnResultFromLocal()); + ctx.addFlightSqlEndpointsLocation(new FlightSqlEndpointsLocation(new TUniqueId(1, 1), + new TNetworkAddress("127.0.0.1", 8070), new TNetworkAddress("127.0.0.1", 8060), new ArrayList<>())); + Assertions.assertEquals(1, ctx.getFlightSqlEndpointsLocations().size()); + // The next statement of the request starts on this frontend again. + adapter.beforeStatement(ctx); + Assertions.assertTrue(ctx.isReturnResultFromLocal()); + // An EXPLAIN is answered here without ever touching a backend, so it never leaves the + // frontend: the sender does not decide where the result is. + ctx.setQueryId(new TUniqueId(2, 2)); + ctx.getResultSender().sendResultSet(resultSet, null, false); + Assertions.assertTrue(ctx.isReturnResultFromLocal()); + Assertions.assertEquals(1, adapter.getChannel().resultNum()); + + // A new request drops everything the previous one left: its deferred coordinator, the + // result nobody pulled, the endpoints, and the result is on this frontend again. + adapter.beforeQuery(ctx); + StmtExecutor deferred = Mockito.mock(StmtExecutor.class); + ctx.addFlightSqlDeferredExecutor(deferred); + adapter.beginRequest(); + Mockito.verify(deferred).finalizeArrowFlightQuery(); + Assertions.assertEquals(0, adapter.getChannel().resultNum()); + Assertions.assertTrue(ctx.getFlightSqlEndpointsLocations().isEmpty()); + Assertions.assertTrue(ctx.isReturnResultFromLocal()); + } + + // A replanned statement is run again without going through beforeStatement. The attempt that + // failed had moved the result to the backends and registered endpoints there; the next one + // starts as the statement did, so that a second failure before its query runs does not leave + // the session believing a result is waiting on the backends, and that nothing the failed + // attempt registered is delivered. + @Test + public void testAnAttemptStartsWhereTheStatementDid() { + ConnectContext ctx = flightSession(); + FlightProtocolAdapter adapter = FlightProtocolAdapter.of(ctx); + + adapter.beforeStatement(ctx); + adapter.beforeAttempt(ctx); + Assertions.assertTrue(ctx.isReturnResultFromLocal()); + Assertions.assertTrue(ctx.getFlightSqlEndpointsLocations().isEmpty()); + + // The first attempt ran its query: result on the backends, one endpoint registered. + adapter.beforeQuery(ctx); + ctx.addFlightSqlEndpointsLocation(new FlightSqlEndpointsLocation(new TUniqueId(1, 1), + new TNetworkAddress("127.0.0.1", 8070), new TNetworkAddress("127.0.0.1", 8060), new ArrayList<>())); + Assertions.assertFalse(ctx.isReturnResultFromLocal()); + Assertions.assertEquals(1, ctx.getFlightSqlEndpointsLocations().size()); + + // It failed and the statement is replanned: the second attempt starts afresh. + adapter.beforeAttempt(ctx); + Assertions.assertTrue(ctx.isReturnResultFromLocal()); + Assertions.assertTrue(ctx.getFlightSqlEndpointsLocations().isEmpty()); + + // Only what the attempt that completes registers is delivered. + adapter.beforeQuery(ctx); + ctx.addFlightSqlEndpointsLocation(new FlightSqlEndpointsLocation(new TUniqueId(2, 2), + new TNetworkAddress("127.0.0.1", 8070), new TNetworkAddress("127.0.0.1", 8060), new ArrayList<>())); + Assertions.assertFalse(ctx.isReturnResultFromLocal()); + Assertions.assertEquals(1, ctx.getFlightSqlEndpointsLocations().size()); + Assertions.assertEquals(new TUniqueId(2, 2), ctx.getFlightSqlEndpointsLocations().get(0).getFinstId()); + + // The next statement of the request leaves that endpoint as it was, whatever its own + // attempts do: a retried attempt withdraws only what the failed one registered. + adapter.beforeStatement(ctx); + adapter.beforeAttempt(ctx); + adapter.beforeQuery(ctx); + ctx.addFlightSqlEndpointsLocation(new FlightSqlEndpointsLocation(new TUniqueId(3, 3), + new TNetworkAddress("127.0.0.1", 8070), new TNetworkAddress("127.0.0.1", 8060), new ArrayList<>())); + adapter.beforeAttempt(ctx); + Assertions.assertEquals(1, ctx.getFlightSqlEndpointsLocations().size()); + Assertions.assertEquals(new TUniqueId(2, 2), ctx.getFlightSqlEndpointsLocations().get(0).getFinstId()); + adapter.beforeQuery(ctx); + ctx.addFlightSqlEndpointsLocation(new FlightSqlEndpointsLocation(new TUniqueId(4, 4), + new TNetworkAddress("127.0.0.1", 8070), new TNetworkAddress("127.0.0.1", 8060), new ArrayList<>())); + Assertions.assertEquals(2, ctx.getFlightSqlEndpointsLocations().size()); + Assertions.assertEquals(new TUniqueId(4, 4), ctx.getFlightSqlEndpointsLocations().get(1).getFinstId()); + + // A new request starts from nothing. + adapter.beginRequest(); + adapter.beforeStatement(ctx); + adapter.beforeAttempt(ctx); + Assertions.assertTrue(ctx.getFlightSqlEndpointsLocations().isEmpty()); + } + @Test public void testOnlyTheLastStatementOfARequestMayReturnAResult() throws Exception { ConnectContext ctx = flightSession(); diff --git a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightResultSenderTest.java b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightResultSenderTest.java index 60e06281386c1d..cb9c2615d96a21 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightResultSenderTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightResultSenderTest.java @@ -67,9 +67,6 @@ public void testResultSetIsCachedUnderTheQueryId() throws Exception { TUniqueId queryId = new TUniqueId(3, 4); ctx.setQueryId(queryId); ctx.setRunningQuery("show variables"); - // The query path marks the result as coming from the backend before an EXPLAIN turns out - // to be answered here; the sender puts that right. - ctx.setReturnResultFromLocal(false); List> rows = Lists.newArrayList(); rows.add(Lists.newArrayList("wait_timeout", "28800")); rows.add(Lists.newArrayList("x", null)); @@ -80,7 +77,6 @@ public void testResultSetIsCachedUnderTheQueryId() throws Exception { ctx.getResultSender().sendResultSet(resultSet, null, false); - Assertions.assertTrue(ctx.isReturnResultFromLocal()); Assertions.assertEquals(1, adapter.getChannel().resultNum()); FlightSqlResultCacheEntry entry = adapter.getChannel().getResult(DebugUtil.printId(queryId)); Assertions.assertNotNull(entry); diff --git a/fe/fe-core/src/test/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapterTest.java b/fe/fe-core/src/test/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapterTest.java index 816ce68ee36c3c..670b33e4f87a13 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapterTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapterTest.java @@ -24,6 +24,7 @@ import org.apache.doris.mysql.DummyMysqlChannel; import org.apache.doris.mysql.MysqlCapability; import org.apache.doris.mysql.MysqlChannel; +import org.apache.doris.mysql.MysqlCommand; import org.apache.doris.mysql.MysqlProto; import org.apache.doris.mysql.MysqlServerStatusFlag; import org.apache.doris.mysql.ProxyMysqlChannel; @@ -32,10 +33,12 @@ import org.apache.doris.qe.ConnectProcessor; import org.apache.doris.qe.ConnectScheduler; import org.apache.doris.qe.protocol.RecordingMysqlChannel; +import org.apache.doris.thrift.TMasterOpRequest; import org.apache.doris.thrift.TResultSinkType; import org.apache.doris.thrift.TUniqueId; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Lists; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; @@ -85,13 +88,33 @@ public void testInternalContextIsAMysqlContextWithoutAClient() { } @Test - public void testProxyContextCollectsThePacketsOfTheForwardedStatement() { + public void testProxyContextCollectsThePacketsOfTheForwardedStatement() throws Exception { ConnectContext ctx = ConnectContext.forMysqlProxy("session-1"); + MysqlProtocolAdapter protocol = MysqlProtocolAdapter.of(ctx); Assertions.assertEquals(ConnectType.MYSQL, ctx.getConnectType()); Assertions.assertTrue(ctx.isProxy()); Assertions.assertEquals("session-1", ctx.getSessionId()); Assertions.assertTrue(ctx.getMysqlChannel() instanceof ProxyMysqlChannel); + + // What the forwarded statement sends is kept, in order, for the frontend the client is + // connected to. + Assertions.assertTrue(protocol.proxyResultPackets().isEmpty()); + ByteBuffer first = ByteBuffer.wrap(new byte[] {1}); + ByteBuffer second = ByteBuffer.wrap(new byte[] {2, 3}); + ctx.getResultSender().sendRow(first); + ctx.getResultSender().sendRow(second); + Assertions.assertEquals(Lists.newArrayList(first, second), protocol.proxyResultPackets()); + + // A failed attempt of the forwarded query may be retried: nothing reached the client, and + // what the attempt wrote is dropped when the next one resets the channel, so the client + // gets the packets of one attempt only. + Assertions.assertTrue(protocol.canRetryQuery(ctx)); + ctx.getResultSender().reset(); + Assertions.assertTrue(protocol.proxyResultPackets().isEmpty()); + ByteBuffer retried = ByteBuffer.wrap(new byte[] {4}); + ctx.getResultSender().sendRow(retried); + Assertions.assertEquals(Lists.newArrayList(retried), protocol.proxyResultPackets()); } @Test @@ -203,6 +226,94 @@ public void testResponsePacketFollowsTheStateAndTheEofCapability() { Assertions.assertTrue(resultSetEnd.remaining() > 5); } + @Test + public void testAMysqlConnectionCanTakeEveryKindOfResult() { + ConnectContext ctx = new ConnectContext(); + MysqlProtocolAdapter protocol = MysqlProtocolAdapter.of(ctx); + + Assertions.assertTrue(protocol.supportsSqlCacheReplay()); + Assertions.assertTrue(protocol.canReplayForwardedQueryResult()); + Assertions.assertTrue(protocol.supportsFeSideResult()); + Assertions.assertTrue(protocol.supportsShortCircuitPointQuery()); + // The result always comes through this frontend, whatever the statement does. + Assertions.assertTrue(ctx.isReturnResultFromLocal()); + protocol.beforeQuery(ctx); + Assertions.assertTrue(ctx.isReturnResultFromLocal()); + + // The FE-side path is taken for a query, not for COM_STMT_EXECUTE. + Assertions.assertTrue(ctx.supportHandleByFe()); + ctx.setCommand(MysqlCommand.COM_STMT_EXECUTE); + Assertions.assertFalse(ctx.supportHandleByFe()); + } + + @Test + public void testAStatementStartsFromAnEmptyChannelAndARetryNeedsOne() throws Exception { + RecordingMysqlChannel channel = new RecordingMysqlChannel(); + ConnectContext ctx = new ConnectContext(new MysqlProtocolAdapter(channel)); + MysqlProtocolAdapter protocol = MysqlProtocolAdapter.of(ctx); + + // What a statement wrote but never flushed is dropped when the next statement starts: + // that is how a request delivers only its last statement's result to a client without + // CLIENT_MULTI_STATEMENTS. + channel.sendOnePacket(ByteBuffer.wrap(new byte[] {1})); + Assertions.assertEquals(1, channel.getOutbound().size()); + protocol.beforeStatement(ctx); + Assertions.assertTrue(channel.getOutbound().isEmpty()); + + // A failed query may be retried while none of its packets reached the client ... + channel.sendOnePacket(ByteBuffer.wrap(new byte[] {2})); + Assertions.assertTrue(protocol.canRetryQuery(ctx)); + // (a replanned attempt starts without touching the channel: the result of every attempt + // is relayed by this frontend, there is nothing of the failed one to withdraw) + protocol.beforeAttempt(ctx); + Assertions.assertEquals(1, channel.getOutbound().size()); + Assertions.assertTrue(protocol.canRetryQuery(ctx)); + // ... and not once one was flushed to the socket. + channel.flush(); + Assertions.assertFalse(protocol.canRetryQuery(ctx)); + // The next statement starts afresh; what went out stays out. + protocol.beforeStatement(ctx); + Assertions.assertTrue(protocol.canRetryQuery(ctx)); + Assertions.assertEquals(1, channel.getOutbound().size()); + } + + @Test + public void testForwardRequestCarriesWhatTheMasterNeedsToAnswerTheClient() { + RecordingMysqlChannel channel = new RecordingMysqlChannel(); + ConnectContext ctx = new ConnectContext(new MysqlProtocolAdapter(channel)); + MysqlProtocolAdapter protocol = MysqlProtocolAdapter.of(ctx); + int flags = MysqlCapability.DEFAULT_CAPABILITY.getFlags() + & ~MysqlCapability.Flag.CLIENT_DEPRECATE_EOF.getFlagBit(); + ctx.setCapability(new MysqlCapability(flags)); + + // A COM_QUERY: the negotiated capabilities, nothing about a prepared statement. + TMasterOpRequest request = new TMasterOpRequest(); + protocol.fillForwardRequest(ctx, request); + Assertions.assertEquals(flags, request.getMysqlCapability()); + Assertions.assertFalse(request.isClientDeprecatedEOF()); + Assertions.assertFalse(request.isSetPrepareExecuteBuffer()); + Assertions.assertFalse(request.isSetCursorFetchRequested()); + + // A COM_STMT_EXECUTE from a client that deprecated EOF: the execute packet and the cursor + // flag travel too. + channel.setClientDeprecatedEOF(); + ctx.setCommand(MysqlCommand.COM_STMT_EXECUTE); + ctx.setPrepareExecuteBuffer(ByteBuffer.wrap(new byte[] {7, 0, 0, 0})); + ctx.setCursorFetchRequested(true); + request = new TMasterOpRequest(); + protocol.fillForwardRequest(ctx, request); + Assertions.assertTrue(request.isClientDeprecatedEOF()); + Assertions.assertArrayEquals(new byte[] {7, 0, 0, 0}, request.getPrepareExecuteBuffer()); + Assertions.assertTrue(request.isCursorFetchRequested()); + + // The master's proxy context takes them over. + ConnectContext proxy = ConnectContext.forMysqlProxy("session-1"); + MysqlProtocolAdapter.of(proxy).restoreFromForwardRequest(proxy, request); + Assertions.assertTrue(proxy.getCapability().isDeprecatedEOF()); + Assertions.assertTrue(proxy.getMysqlChannel().clientDeprecatedEOF()); + Assertions.assertTrue(proxy.isCursorFetchRequested()); + } + @Test public void testConnectionLifecycleGoesThroughTheChannel() throws Exception { MysqlChannel channel = Mockito.mock(MysqlChannel.class); diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/FEOpExecutorMysqlProtocolTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/FEOpExecutorMysqlProtocolTest.java index b8a441c7e5650d..f3a019badb3647 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/FEOpExecutorMysqlProtocolTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/FEOpExecutorMysqlProtocolTest.java @@ -23,6 +23,7 @@ import org.apache.doris.mysql.MysqlCapability; import org.apache.doris.mysql.MysqlCommand; import org.apache.doris.mysql.MysqlProto; +import org.apache.doris.mysql.protocol.MysqlProtocolAdapter; import org.apache.doris.system.SystemInfoService; import org.apache.doris.thrift.TMasterOpRequest; import org.apache.doris.thrift.TMasterOpResult; @@ -113,7 +114,7 @@ public void testForwardedCapabilityAndMissingCursorFlag() throws Exception { TMasterOpRequest request = new TMasterOpRequest(); request.setMysqlCapability(legacyFlags); ConnectContext context = createContext(); - ConnectProcessor.restoreForwardedMysqlContext(context, request); + MysqlProtocolAdapter.of(context).restoreFromForwardRequest(context, request); Assertions.assertEquals(legacyFlags, context.getCapability().getFlags()); Assertions.assertFalse(context.getMysqlChannel().getSerializer().getCapability().isDeprecatedEOF()); @@ -122,14 +123,14 @@ public void testForwardedCapabilityAndMissingCursorFlag() throws Exception { request.setPrepareExecuteBuffer(new byte[] {0}); context = createContext(); context.setConnectAttributes(ImmutableMap.of("_client_name", "MySQL Connector/J", "_client_version", "8.2.0")); - ConnectProcessor.restoreForwardedMysqlContext(context, request); + MysqlProtocolAdapter.of(context).restoreFromForwardRequest(context, request); Assertions.assertFalse(context.isCursorFetchRequested()); Assertions.assertTrue(context.getCapability().isDeprecatedEOF()); request.setCursorFetchRequested(true); - ConnectProcessor.restoreForwardedMysqlContext(context, request); + MysqlProtocolAdapter.of(context).restoreFromForwardRequest(context, request); Assertions.assertTrue(context.isCursorFetchRequested()); request.setCursorFetchRequested(false); - ConnectProcessor.restoreForwardedMysqlContext(context, request); + MysqlProtocolAdapter.of(context).restoreFromForwardRequest(context, request); Assertions.assertFalse(context.isCursorFetchRequested()); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/ProtocolCapabilityWiringTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/ProtocolCapabilityWiringTest.java new file mode 100644 index 00000000000000..21898d0f1a623d --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/ProtocolCapabilityWiringTest.java @@ -0,0 +1,317 @@ +// 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.qe; + +import org.apache.doris.analysis.UserIdentity; +import org.apache.doris.arrowflight.FlightSqlConnectProcessor; +import org.apache.doris.arrowflight.protocol.FlightProtocolAdapter; +import org.apache.doris.catalog.Env; +import org.apache.doris.common.AnalysisException; +import org.apache.doris.common.Config; +import org.apache.doris.common.UserException; +import org.apache.doris.common.jmockit.Deencapsulation; +import org.apache.doris.ha.FrontendNodeType; +import org.apache.doris.mysql.MysqlCapability; +import org.apache.doris.mysql.MysqlProto; +import org.apache.doris.mysql.protocol.MysqlProtocolAdapter; +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.plans.logical.LogicalPlan; +import org.apache.doris.qe.QueryState.MysqlStateType; +import org.apache.doris.qe.protocol.RecordingMysqlChannel; +import org.apache.doris.qe.protocol.RecordingMysqlChannel.RecordedPacket; +import org.apache.doris.rpc.RpcException; +import org.apache.doris.system.SystemInfoService; +import org.apache.doris.thrift.TMasterOpRequest; +import org.apache.doris.thrift.TMasterOpResult; +import org.apache.doris.thrift.TUniqueId; +import org.apache.doris.utframe.MockedBackendFactory.DefaultPBackendServiceImpl; +import org.apache.doris.utframe.TestWithFeService; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.nio.ByteBuffer; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * The execution layer does not branch on the protocol a connection speaks; it asks the + * connection's {@code ProtocolAdapter}. The adapter tests pin what each adapter answers; these + * tests drive real statements through the processors and the executor and check that the answer + * is what decides, where a unit test of the adapter alone cannot show it. + */ +public class ProtocolCapabilityWiringTest extends TestWithFeService { + private static final String DB_NAME = "protocol_capability_wiring_db"; + + @Override + protected void runBeforeAll() throws Exception { + createDatabaseAndUse(DB_NAME); + } + + // On a follower a query is forwarded to the master when the session asks for it + // (ForceForwardAllQueriesTest); an Arrow Flight SQL session cannot take the master's answer + // and is refused before any rpc, a MySQL connection goes on to forward. + @Test + public void testForwardedQueryIsRefusedOnAFlightSessionBeforeTheRpc() throws Exception { + Env env = Env.getCurrentEnv(); + FrontendNodeType originalFeType = env.getFeType(); + AtomicBoolean canRead = Deencapsulation.getField(env, "canRead"); + boolean originalCanRead = canRead.get(); + boolean originalForceForward = Config.force_forward_all_queries; + Deencapsulation.setField(env, "feType", FrontendNodeType.FOLLOWER); + canRead.set(true); + Config.force_forward_all_queries = false; + try { + ConnectContext flight = flightContext(); + flight.getSessionVariable().forceForwardAllQueries = true; + StmtExecutor refused = new StmtExecutor(flight, analyzeAndGetStmtByNereids("select 1", flight)); + + UserException e = Assertions.assertThrows(UserException.class, refused::execute); + + Assertions.assertTrue(e.getMessage().contains("not supported on an Arrow Flight SQL connection"), + e.getMessage()); + Assertions.assertFalse(refused.hasForwardedToMaster()); + Assertions.assertEquals(MysqlStateType.ERR, flight.getState().getStateType()); + + ConnectContext mysql = createDefaultCtx(); + mysql.setDatabase(DB_NAME); + mysql.setThreadLocalInfo(); + mysql.getSessionVariable().forceForwardAllQueries = true; + StmtExecutor forwarded = new StmtExecutor(mysql, analyzeAndGetStmtByNereids("select 1", mysql)); + try { + forwarded.execute(); + } catch (Exception expected) { + // There is no master to reach from a unit test; what matters is that the + // statement was handed to the forwarding path instead of being refused. + } + Assertions.assertTrue(forwarded.hasForwardedToMaster()); + } finally { + Deencapsulation.setField(env, "feType", originalFeType); + canRead.set(originalCanRead); + Config.force_forward_all_queries = originalForceForward; + connectContext.setThreadLocalInfo(); + } + } + + // The previous statement of a Flight request left its result on a backend; the next one is + // answered by this frontend again, because the processor tells the adapter where a statement + // starts. + @Test + public void testAStatementStartsOnTheFrontendWhateverThePreviousOneDid() throws Exception { + ConnectContext flight = flightContext(); + FlightProtocolAdapter adapter = FlightProtocolAdapter.of(flight); + adapter.beforeQuery(flight); + Assertions.assertFalse(flight.isReturnResultFromLocal()); + try (FlightSqlConnectProcessor processor = new FlightSqlConnectProcessor(flight)) { + processor.handleQuery("show variables like 'wait_timeout'"); + + Assertions.assertNotEquals(MysqlStateType.ERR, flight.getState().getStateType(), + flight.getState().getErrorMessage()); + Assertions.assertTrue(flight.isReturnResultFromLocal()); + Assertions.assertEquals(1, adapter.getChannel().resultNum()); + } finally { + adapter.getChannel().close(); + connectContext.setThreadLocalInfo(); + } + } + + // The master answers a forwarded query for the client of the frontend that forwarded it: the + // packets it collects and the response that ends them follow the capabilities the request + // carries, restored into the proxy context by the adapter. + @Test + public void testMasterAnswersAForwardedQueryWithTheClientsCapabilities() throws Exception { + for (boolean deprecateEof : new boolean[] {true, false}) { + int flags = MysqlCapability.DEFAULT_CAPABILITY.getFlags(); + if (!deprecateEof) { + flags &= ~MysqlCapability.Flag.CLIENT_DEPRECATE_EOF.getFlagBit(); + } + TMasterOpRequest request = new TMasterOpRequest(); + request.setDb(DB_NAME); + request.setUser("root"); + request.setCurrentUserIdent(UserIdentity.ROOT.toThrift()); + request.setSql("select 1"); + request.setMysqlCapability(flags); + request.setClientDeprecatedEOF(deprecateEof); + ConnectContext proxy = ConnectContext.forMysqlProxy("session-1"); + try { + TMasterOpResult result = new MysqlConnectProcessor(proxy).proxyExecute(request); + + // A query's state is EOF, the end of its result set; only an OK state sets statusCode 0. + Assertions.assertEquals(MysqlStateType.EOF.name(), result.getStatus(), + proxy.getState().getErrorMessage()); + // column count, one column definition, the terminator the client expects, one row + List packets = result.getQueryResultBufList(); + Assertions.assertEquals(deprecateEof ? 3 : 4, packets.size()); + Assertions.assertEquals(1, MysqlProto.readVInt(packets.get(0).duplicate())); + if (!deprecateEof) { + ByteBuffer eof = packets.get(2); + Assertions.assertEquals(0xFE, Byte.toUnsignedInt(eof.get(eof.position()))); + Assertions.assertEquals(5, eof.remaining()); + } + // and the packet that ends the result set is an EOF, or an OK with the 0xFE header + ByteBuffer end = result.packet; + Assertions.assertEquals(0xFE, Byte.toUnsignedInt(end.get(end.position()))); + Assertions.assertEquals(deprecateEof, end.remaining() > 5); + Assertions.assertEquals(deprecateEof, result.isClientDeprecatedEofApplied()); + } finally { + connectContext.setThreadLocalInfo(); + } + } + } + + // A query whose backend rpc failed is run again under a new query id on a MySQL connection: its + // client has seen nothing of the failed attempt. A Flight session gets the failure instead, + // because the backend endpoints the failed attempt registered would have to be withdrawn + // first. The mocked backend answers the first exec_plan_fragment rpc with a TIMEOUT status, + // which the coordinator raises as the RpcException the executor retries on. + @Test + public void testAFailedQueryIsRetriedOnlyWhereTheClientCannotTell() throws Exception { + createTable("create table retry_tbl (k int) distributed by hash(k) buckets 1" + + " properties ('replication_num' = '1')"); + try { + ConnectContext mysql = createDefaultCtx(); + mysql.setDatabase(DB_NAME); + mysql.setThreadLocalInfo(); + // The table is empty; without this the scan folds into an empty relation the frontend + // answers by itself and no backend rpc is ever sent. + mysql.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION"); + StmtExecutor retried = executorFor(mysql, "select k from retry_tbl"); + TUniqueId firstQueryId = new TUniqueId(1, 1); + int callsBefore = DefaultPBackendServiceImpl.getExecPlanFragmentCalls(); + DefaultPBackendServiceImpl.failNextExecPlanFragments(1); + + retried.execute(firstQueryId); + + Assertions.assertEquals(MysqlStateType.EOF, mysql.getState().getStateType(), + mysql.getState().getErrorMessage()); + // The rpc that failed, then at least the one that succeeded, under a new query id. + Assertions.assertEquals(0, DefaultPBackendServiceImpl.getPendingExecPlanFragmentFailures()); + Assertions.assertTrue(DefaultPBackendServiceImpl.getExecPlanFragmentCalls() - callsBefore >= 2); + Assertions.assertNotEquals(firstQueryId, mysql.queryId()); + + ConnectContext flight = flightContext(); + flight.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION"); + StmtExecutor failed = executorFor(flight, "select k from retry_tbl"); + TUniqueId flightQueryId = new TUniqueId(2, 2); + DefaultPBackendServiceImpl.failNextExecPlanFragments(1); + + RpcException e = Assertions.assertThrows(RpcException.class, () -> failed.execute(flightQueryId)); + + Assertions.assertTrue(e.getMessage().contains("injected exec_plan_fragment timeout"), e.getMessage()); + Assertions.assertEquals(0, DefaultPBackendServiceImpl.getPendingExecPlanFragmentFailures()); + Assertions.assertEquals(MysqlStateType.ERR, flight.getState().getStateType()); + Assertions.assertEquals(flightQueryId, flight.queryId()); + } finally { + DefaultPBackendServiceImpl.failNextExecPlanFragments(0); + connectContext.setThreadLocalInfo(); + dropTable("retry_tbl", true); + } + } + + // A query that fails with an error the statement is replanned on is run again by queryRetry + // without going back through the processor: it is the executor that starts every attempt + // where the statement started. The first attempt of this Flight query moved its result to + // the backend and registered an endpoint there before it failed; the second fails while it + // is planned, before it could run a query. The session must not be left believing a result + // is waiting on the backend -- that state would keep its cleanup (StatementContext.close, + // the query registration) waiting for a DoGet that never comes -- and the endpoint the failed + // attempt registered must not be delivered. + @Test + public void testEveryAttemptOfAReplannedFlightQueryStartsOnTheFrontend() throws Exception { + createTable("create table replan_tbl (k int) distributed by hash(k) buckets 1" + + " properties ('replication_num' = '1')"); + try { + ConnectContext flight = flightContext(); + flight.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION"); + StmtExecutor executor = Mockito.spy(executorFor(flight, "select k from replan_tbl")); + flight.setExecutor(executor); + // The second attempt fails before its query is run. + Mockito.doCallRealMethod().doThrow(new AnalysisException("injected planning failure")) + .when(executor).checkBlockRulesByRegex(Mockito.any()); + DefaultPBackendServiceImpl.failNextExecPlanFragments(1, + SystemInfoService.ERROR_E230 + " injected replan error"); + + executor.execute(); + + Assertions.assertEquals(0, DefaultPBackendServiceImpl.getPendingExecPlanFragmentFailures()); + Mockito.verify(executor, Mockito.times(2)).checkBlockRulesByRegex(Mockito.any()); + Assertions.assertEquals(MysqlStateType.ERR, flight.getState().getStateType()); + Assertions.assertTrue(flight.getState().getErrorMessage().contains("injected planning failure"), + flight.getState().getErrorMessage()); + Assertions.assertTrue(flight.isReturnResultFromLocal()); + Assertions.assertTrue(flight.getFlightSqlEndpointsLocations().isEmpty()); + } finally { + DefaultPBackendServiceImpl.failNextExecPlanFragments(0); + connectContext.setThreadLocalInfo(); + dropTable("replan_tbl", true); + } + } + + // An internal executor (the dry run of RefreshMTMVCommand) answers the client of the session + // that issued the statement, encoded with that client's capabilities: it is handed the + // caller's sender, as its own session is connected to nobody. + @Test + public void testAnInternalExecutorAnswersOnTheCallersConnection() throws Exception { + // The caller's client did not deprecate EOF; an internal session's default capabilities do. + RecordingMysqlChannel callerChannel = new RecordingMysqlChannel(); + ConnectContext caller = new ConnectContext(new MysqlProtocolAdapter(callerChannel)); + ConnectContext internal = createDefaultCtx(); + internal.setDatabase(DB_NAME); + internal.setThreadLocalInfo(); + try { + StmtExecutor executor = executorFor(internal, "select 1"); + + executor.executeInternalQueryAndSend((LogicalPlanAdapter) executor.getParsedStmt(), + caller.getResultSender()); + + // column count, one column definition, the terminator the caller's client expects, one row + List packets = callerChannel.getOutbound(); + Assertions.assertEquals(4, packets.size()); + Assertions.assertArrayEquals(new byte[] {1}, packets.get(0).getPayload()); + Assertions.assertEquals(0xFE, Byte.toUnsignedInt(packets.get(2).getPayload()[0])); + Assertions.assertEquals(5, packets.get(2).getPayload().length); + Assertions.assertArrayEquals(new byte[] {1, '1'}, packets.get(3).getPayload()); + } finally { + connectContext.setThreadLocalInfo(); + } + } + + // An executor the way the connect processor builds one: the parsed statement carries the + // original text, and the statement context knows it (so a NereidsCoordinator runs the query). + private static StmtExecutor executorFor(ConnectContext ctx, String sql) { + StatementContext statementContext = new StatementContext(ctx, new OriginStatement(sql, 0)); + ctx.setStatementContext(statementContext); + LogicalPlan plan = new NereidsParser().parseSingle(sql); + LogicalPlanAdapter adapter = new LogicalPlanAdapter(plan, statementContext); + adapter.setOrigStmt(statementContext.getOriginStatement()); + return new StmtExecutor(ctx, adapter); + } + + private ConnectContext flightContext() { + ConnectContext ctx = ConnectContext.forFlight("test-peer-identity"); + ctx.setCurrentUserIdentity(UserIdentity.ROOT); + ctx.setRemoteIP("127.0.0.1"); + ctx.setEnv(Env.getCurrentEnv()); + ctx.setDatabase(DB_NAME); + ctx.setThreadLocalInfo(); + return ctx; + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/protocol/MysqlPacketGoldenTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/protocol/MysqlPacketGoldenTest.java index f3d6c7037ef832..871f99b310f33f 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/protocol/MysqlPacketGoldenTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/protocol/MysqlPacketGoldenTest.java @@ -59,7 +59,8 @@ *
  • errors: a syntax error, an unknown table, and an error followed by a healthy statement on * the same connection;
    • *
  • multi-statement requests with and without {@code CLIENT_MULTI_STATEMENTS}, which decides - * whether the intermediate result set gets a terminator at all;
    • + * whether the intermediate result set gets a terminator at all, ending in a query, in a + * {@code SET} and in a {@code BEGIN}; *
  • connection commands: {@code COM_FIELD_LIST}, {@code COM_STMT_PREPARE}, * {@code COM_STMT_CLOSE}, {@code COM_SET_OPTION}, {@code COM_RESET_CONNECTION}, * {@code COM_PING}, {@code COM_INIT_DB}, {@code COM_STATISTICS}, an unknown command, and @@ -139,6 +140,23 @@ private List cases() { .add(query("select 1; select 2"))); cases.add(new GoldenCase("multi-statement-without-capability", MODERN_CLIENT) .add(query("select 1; select 2"))); + // The same two ways of finishing a request whose last statement is not a query. + cases.add(new GoldenCase("multi-statement-with-capability-query-then-set", MULTI_STATEMENT_CLIENT) + .add(query("select 1; set @a = 1"))); + cases.add(new GoldenCase("multi-statement-without-capability-query-then-set", MODERN_CLIENT) + .add(query("select 1; set @a = 1"))); + // A later statement of the request fails. With the capability the client already got the + // first statement's result; without it, the client gets only the ERR, numbered from where + // the last flush left off -- nothing of this request had reached it yet. + cases.add(new GoldenCase("multi-statement-with-capability-query-then-error", MULTI_STATEMENT_CLIENT) + .add(query("select 1; select * from no_such_table"))); + cases.add(new GoldenCase("multi-statement-without-capability-query-then-error", MODERN_CLIENT) + .add(query("select 1; select * from no_such_table"))); + // A transaction command used to reset the channel on its own; only the shape is kept, the + // OK carries a label derived from the query id. + cases.add(new GoldenCase("multi-statement-without-capability-query-then-begin", MODERN_CLIENT, + ProtocolGolden.Fidelity.SUMMARY) + .add(query("select 1; begin"))); cases.add(new GoldenCase("com-field-list", MODERN_CLIENT) .add(fieldList(TABLE_NAME))); cases.add(new GoldenCase("com-stmt-prepare", MODERN_CLIENT) diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/protocol/RecordingMysqlChannel.java b/fe/fe-core/src/test/java/org/apache/doris/qe/protocol/RecordingMysqlChannel.java index 457d29b329872a..763e7734868611 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/protocol/RecordingMysqlChannel.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/protocol/RecordingMysqlChannel.java @@ -34,6 +34,12 @@ * the payload, then advances the sequence id. There is no send buffer here, so the header is not * materialized; the sequence id each packet would have carried is recorded next to the payload * instead, and the golden renders the header from the two. + * + *

    What the send buffer does is modeled, though: a flush pushes everything written so far, and + * {@link #reset()} drops what was written after the last flush and rewinds the sequence id to the + * last one the client saw, the way {@link org.apache.doris.mysql.MysqlChannel#reset()} does. So + * the golden shows what reaches the client, not everything the server wrote, numbered as the + * client will see it. */ public class RecordingMysqlChannel extends DummyMysqlChannel { @@ -105,6 +111,7 @@ public ByteBuffer fetchOnePacket() { // The real channel advances the sequence id once per packet it reads, so the first response // packet is framed with the request's id plus one. accSequenceId(); + wireSequenceId = sequenceId; return packet; } @@ -116,13 +123,31 @@ public void sendOnePacket(ByteBuffer packet) { @Override public void sendAndFlush(ByteBuffer packet) { record(packet, true); + wireSequenceId = sequenceId; + isSend = true; } @Override public void flush() { - if (!outbound.isEmpty()) { + // Nothing to push when the last packet already went out with a flush. + if (!outbound.isEmpty() && !outbound.get(outbound.size() - 1).isFlushed()) { outbound.get(outbound.size() - 1).markFlushed(); + wireSequenceId = sequenceId; + isSend = true; + } + } + + @Override + public void reset() { + isSend = false; + // A flush pushes everything written so far; what was written after the last one is still + // in the buffer, and that is what a reset throws away. + int end = outbound.size(); + while (end > 0 && !outbound.get(end - 1).isFlushed()) { + end--; } + outbound.subList(end, outbound.size()).clear(); + sequenceId = wireSequenceId; } private void record(ByteBuffer packet, boolean flushed) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/utframe/MockedBackendFactory.java b/fe/fe-core/src/test/java/org/apache/doris/utframe/MockedBackendFactory.java index 5c20a2b64f6b90..163ceaf4ae9188 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/utframe/MockedBackendFactory.java +++ b/fe/fe-core/src/test/java/org/apache/doris/utframe/MockedBackendFactory.java @@ -92,6 +92,7 @@ import java.util.List; import java.util.Random; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; /* @@ -494,13 +495,61 @@ public java.util.List getPythonPackages(String pythonVersion // The default Brpc service. public static class DefaultPBackendServiceImpl extends PBackendServiceGrpc.PBackendServiceImplBase { + // How many of the next exec_plan_fragment rpcs answer with a failed status, without + // blacklisting this backend. Every exec_plan_fragment rpc counts. By default the status is + // a TIMEOUT, which the coordinator reports as an RpcException (the kind of failure + // StmtExecutor retries a query on); with an error message it is an INTERNAL_ERROR carrying + // that message, which the coordinator reports as a UserException. + private static final AtomicInteger execPlanFragmentTimeouts = new AtomicInteger(); + private static volatile String execPlanFragmentErrorMsg = null; + private static final AtomicInteger execPlanFragmentCalls = new AtomicInteger(); + + /** Makes the next {@code times} exec_plan_fragment rpcs of every mocked backend time out. */ + public static void failNextExecPlanFragments(int times) { + failNextExecPlanFragments(times, null); + } + + /** + * Makes the next {@code times} exec_plan_fragment rpcs of every mocked backend fail with + * {@code errorMsg} (a TIMEOUT when it is null), e.g. with one of + * {@code SystemInfoService.NEED_REPLAN_ERRORS} to have the statement replanned. + */ + public static void failNextExecPlanFragments(int times, String errorMsg) { + execPlanFragmentErrorMsg = errorMsg; + execPlanFragmentTimeouts.set(times); + } + + public static int getExecPlanFragmentCalls() { + return execPlanFragmentCalls.get(); + } + + /** How many of the failures asked for by {@link #failNextExecPlanFragments} are still pending. */ + public static int getPendingExecPlanFragmentFailures() { + return execPlanFragmentTimeouts.get(); + } + + private static InternalService.PExecPlanFragmentResult execPlanFragmentResult() { + execPlanFragmentCalls.incrementAndGet(); + if (execPlanFragmentTimeouts.getAndUpdate(left -> left > 0 ? left - 1 : 0) > 0) { + String errorMsg = execPlanFragmentErrorMsg; + if (errorMsg == null) { + return InternalService.PExecPlanFragmentResult.newBuilder() + .setStatus(Types.PStatus.newBuilder().setStatusCode(TStatusCode.TIMEOUT.getValue()) + .addErrorMsgs("injected exec_plan_fragment timeout")).build(); + } + return InternalService.PExecPlanFragmentResult.newBuilder() + .setStatus(Types.PStatus.newBuilder().setStatusCode(TStatusCode.INTERNAL_ERROR.getValue()) + .addErrorMsgs(errorMsg)).build(); + } + return InternalService.PExecPlanFragmentResult.newBuilder() + .setStatus(Types.PStatus.newBuilder().setStatusCode(0)).build(); + } @Override public void execPlanFragment(InternalService.PExecPlanFragmentRequest request, StreamObserver responseObserver) { System.out.println("get exec_plan_fragment request"); - responseObserver.onNext(InternalService.PExecPlanFragmentResult.newBuilder() - .setStatus(Types.PStatus.newBuilder().setStatusCode(0)).build()); + responseObserver.onNext(execPlanFragmentResult()); responseObserver.onCompleted(); } @@ -508,8 +557,7 @@ public void execPlanFragment(InternalService.PExecPlanFragmentRequest request, public void execPlanFragmentPrepare(InternalService.PExecPlanFragmentRequest request, StreamObserver responseObserver) { System.out.println("get exec_plan_fragment_prepare request"); - responseObserver.onNext(InternalService.PExecPlanFragmentResult.newBuilder() - .setStatus(Types.PStatus.newBuilder().setStatusCode(0)).build()); + responseObserver.onNext(execPlanFragmentResult()); responseObserver.onCompleted(); } @@ -517,8 +565,7 @@ public void execPlanFragmentPrepare(InternalService.PExecPlanFragmentRequest req public void execPlanFragmentStart(InternalService.PExecPlanFragmentStartRequest request, StreamObserver responseObserver) { System.out.println("get exec_plan_fragment_start request"); - responseObserver.onNext(InternalService.PExecPlanFragmentResult.newBuilder() - .setStatus(Types.PStatus.newBuilder().setStatusCode(0)).build()); + responseObserver.onNext(execPlanFragmentResult()); responseObserver.onCompleted(); } diff --git a/fe/fe-core/src/test/resources/protocol-golden/mysql-packets.txt b/fe/fe-core/src/test/resources/protocol-golden/mysql-packets.txt index d22f48b3c82954..282beb41fc53c5 100644 --- a/fe/fe-core/src/test/resources/protocol-golden/mysql-packets.txt +++ b/fe/fe-core/src/test/resources/protocol-golden/mysql-packets.txt @@ -258,6 +258,19 @@ client capability: deprecate_eof=true multi_statements=true === case multi-statement-without-capability === client capability: deprecate_eof=true multi_statements=false --> COM_QUERY select 1; select 2 + packet seq=1 len=1 kind=PAYLOAD + 0000 01 |.| + packet seq=2 len=23 kind=PAYLOAD + 0000 03 64 65 66 00 00 00 01 32 00 0c 21 00 04 00 00 |.def....2..!....| + 0010 00 01 00 00 00 00 00 |.......| + packet seq=3 len=2 kind=PAYLOAD + 0000 01 32 |.2| + packet seq=4 len=8 kind=EOF flushed + 0000 fe 00 00 00 00 00 00 00 |........| + +=== case multi-statement-with-capability-query-then-set === +client capability: deprecate_eof=true multi_statements=true +--> COM_QUERY select 1; set @a = 1 packet seq=1 len=1 kind=PAYLOAD 0000 01 |.| packet seq=2 len=23 kind=PAYLOAD @@ -265,15 +278,56 @@ client capability: deprecate_eof=true multi_statements=false 0010 00 01 00 00 00 00 00 |.......| packet seq=3 len=2 kind=PAYLOAD 0000 01 31 |.1| - packet seq=4 len=1 kind=PAYLOAD + packet seq=4 len=8 kind=EOF flushed + 0000 fe 00 00 08 00 00 00 00 |........| + packet seq=5 len=8 kind=OK flushed + 0000 00 00 00 00 00 00 00 00 |........| + +=== case multi-statement-without-capability-query-then-set === +client capability: deprecate_eof=true multi_statements=false +--> COM_QUERY select 1; set @a = 1 + packet seq=1 len=8 kind=OK flushed + 0000 00 00 00 00 00 00 00 00 |........| + +=== case multi-statement-with-capability-query-then-error === +client capability: deprecate_eof=true multi_statements=true +--> COM_QUERY select 1; select * from no_such_table + packet seq=1 len=1 kind=PAYLOAD 0000 01 |.| - packet seq=5 len=23 kind=PAYLOAD - 0000 03 64 65 66 00 00 00 01 32 00 0c 21 00 04 00 00 |.def....2..!....| + packet seq=2 len=23 kind=PAYLOAD + 0000 03 64 65 66 00 00 00 01 31 00 0c 21 00 04 00 00 |.def....1..!....| 0010 00 01 00 00 00 00 00 |.......| - packet seq=6 len=2 kind=PAYLOAD - 0000 01 32 |.2| - packet seq=7 len=8 kind=EOF flushed - 0000 fe 00 00 00 00 00 00 00 |........| + packet seq=3 len=2 kind=PAYLOAD + 0000 01 31 |.1| + packet seq=4 len=8 kind=EOF flushed + 0000 fe 00 00 08 00 00 00 00 |........| + packet seq=5 len=124 kind=ERR flushed + 0000 ff 51 04 23 48 59 30 30 30 65 72 72 43 6f 64 65 |.Q.#HY000errCode| + 0010 20 3d 20 32 2c 20 64 65 74 61 69 6c 4d 65 73 73 | = 2, detailMess| + 0020 61 67 65 20 3d 20 54 61 62 6c 65 20 5b 6e 6f 5f |age = Table [no_| + 0030 73 75 63 68 5f 74 61 62 6c 65 5d 20 64 6f 65 73 |such_table] does| + 0040 20 6e 6f 74 20 65 78 69 73 74 20 69 6e 20 64 61 | not exist in da| + 0050 74 61 62 61 73 65 20 5b 70 72 6f 74 6f 63 6f 6c |tabase [protocol| + 0060 5f 67 6f 6c 64 65 6e 5f 64 62 5d 2e 28 6c 69 6e |_golden_db].(lin| + 0070 65 20 31 2c 20 70 6f 73 20 32 34 29 |e 1, pos 24)| + +=== case multi-statement-without-capability-query-then-error === +client capability: deprecate_eof=true multi_statements=false +--> COM_QUERY select 1; select * from no_such_table + packet seq=1 len=124 kind=ERR flushed + 0000 ff 51 04 23 48 59 30 30 30 65 72 72 43 6f 64 65 |.Q.#HY000errCode| + 0010 20 3d 20 32 2c 20 64 65 74 61 69 6c 4d 65 73 73 | = 2, detailMess| + 0020 61 67 65 20 3d 20 54 61 62 6c 65 20 5b 6e 6f 5f |age = Table [no_| + 0030 73 75 63 68 5f 74 61 62 6c 65 5d 20 64 6f 65 73 |such_table] does| + 0040 20 6e 6f 74 20 65 78 69 73 74 20 69 6e 20 64 61 | not exist in da| + 0050 74 61 62 61 73 65 20 5b 70 72 6f 74 6f 63 6f 6c |tabase [protocol| + 0060 5f 67 6f 6c 64 65 6e 5f 64 62 5d 2e 28 6c 69 6e |_golden_db].(lin| + 0070 65 20 31 2c 20 70 6f 73 20 32 34 29 |e 1, pos 24)| + +=== case multi-statement-without-capability-query-then-begin === +client capability: deprecate_eof=true multi_statements=false +--> COM_QUERY select 1; begin + kinds: OK flushed === case com-field-list === client capability: deprecate_eof=true multi_statements=false diff --git a/regression-test/suites/arrow_flight_sql_p0/test_arrow_flight_session_lifecycle.groovy b/regression-test/suites/arrow_flight_sql_p0/test_arrow_flight_session_lifecycle.groovy new file mode 100644 index 00000000000000..035ae358ae4b2e --- /dev/null +++ b/regression-test/suites/arrow_flight_sql_p0/test_arrow_flight_session_lifecycle.groovy @@ -0,0 +1,143 @@ +// 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.doris.regression.util.JdbcUtils + +import java.sql.Types + +// One Arrow Flight SQL session, many requests. A statement the frontend answers itself (SHOW, SET, +// EXPLAIN, DESC, DDL) has its result cached on the session for the client's DoGet, every column as +// text; a query is run on the backend and pulled from there, typed; underneath, the session core +// (variables, current database) is one and the same. The session's protocol adapter decides per +// statement where the result is (FlightProtocolAdapter.beforeStatement / beforeQuery) and drops +// what the previous request left behind (beginRequest), so a session can alternate between the two +// kinds of statements without one leaking into the next. +// +// Not in the 'arrow_flight_sql' group on purpose: `sql` stays the MySQL control connection, and the +// statements under test go to the raw Flight connection so that nothing prepends `USE ;`. +suite("test_arrow_flight_session_lifecycle") { + def tableName = "arrow_flight_session_lifecycle" + sql "DROP TABLE IF EXISTS ${tableName}" + sql """ + CREATE TABLE ${tableName} (k INT, v VARCHAR(20)) + DUPLICATE KEY(k) DISTRIBUTED BY HASH(k) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + sql "INSERT INTO ${tableName} VALUES (1, 'a'), (2, 'b'), (3, 'c')" + + def flightConn = context.getArrowFlightSqlConnection() + + // The rows a statement returns over the Flight session, and the JDBC types of its columns. + def flight = { String stmt -> + logger.info("flight: ${stmt}".toString()) + def (rows, meta) = JdbcUtils.executeQueryToList(flightConn, stmt) + def types = (1..meta.getColumnCount()).collect { meta.getColumnType(it) } + def names = (1..meta.getColumnCount()).collect { meta.getColumnName(it) } + return [rows: rows, types: types, names: names] + } + // Arrow JDBC hands a TINYINT back as a Byte and a BIGINT as a Long; compare the values, not the boxes. + def nums = { List> rowList -> rowList.collect { row -> row.collect { it instanceof Number ? ((Number) it).longValue() : it } } } + def flightFails = { String stmt, String fragment -> + try { + flight(stmt) + } catch (Exception e) { + assertTrue(e.getMessage().contains(fragment), + "expected the error of '${stmt}' to mention '${fragment}', got: ${e.getMessage()}") + return + } + throw new AssertionError("'${stmt}' did not fail with '${fragment}'") + } + + // The session's current database is set by a request of its own and holds for the rest. + def use = flight("USE ${context.dbName}") + assertEquals(["StatusResult"], use.names) + + // 1. A result the frontend materializes itself is delivered from the frontend, as text. + def variables = flight("SHOW VARIABLES LIKE 'wait_timeout'") + assertEquals(1, variables.rows.size()) + assertEquals("wait_timeout", variables.rows[0][0].toString()) + assertTrue(variables.types.every { it == Types.VARCHAR }, "expected text columns, got ${variables.types}") + + // 2. A query is run on the backend and pulled from there, typed. That includes a query the + // planner could answer on the frontend: a Flight client expects typed Arrow data, and the + // frontend-side result is text (FlightProtocolAdapter.supportsFeSideResult). + def rows = flight("SELECT k, v FROM ${tableName} ORDER BY k") + assertEquals([[1L, "a"], [2L, "b"], [3L, "c"]], nums(rows.rows)) + assertEquals([Types.INTEGER, Types.VARCHAR], rows.types) + def literal = flight("SELECT 1") + assertEquals([[1L]], nums(literal.rows)) + assertEquals(Types.TINYINT, literal.types[0]) + def sessionVar = flight("SELECT @@wait_timeout") + assertTrue(sessionVar.types[0] in [Types.INTEGER, Types.BIGINT], "session variable came back as ${sessionVar.types}") + + // 3. The two kinds of statements alternate on one session; each request delivers exactly the + // result of its own statement, wherever the previous one left its result. + assertEquals([[1L]], nums(flight("SELECT 1").rows)) + assertEquals("wait_timeout", flight("SHOW VARIABLES LIKE 'wait_timeout'").rows[0][0].toString()) + assertEquals([[2L]], nums(flight("SELECT k FROM ${tableName} WHERE k = 2").rows)) + flight("SET enable_profile = false") + assertEquals([[3L]], nums(flight("SELECT count(*) FROM ${tableName}").rows)) + def explain = flight("EXPLAIN SELECT k FROM ${tableName}") + assertTrue(explain.rows.collect { it[0].toString() }.join("\n").contains("OlapScanNode"), + "EXPLAIN over Flight returned no plan: ${explain.rows}") + assertEquals([Types.VARCHAR], explain.types) + def desc = flight("DESC ${tableName}") + assertEquals(["k", "v"], desc.rows.collect { it[0].toString() }) + assertEquals([[3L]], nums(flight("SELECT max(k) FROM ${tableName}").rows)) + + // 4. EXPLAIN PLAN PROCESS has a result over Flight too; it used to return nothing because it + // bypassed the frontend-side result path. + def process = flight("EXPLAIN PLAN PROCESS SELECT k FROM ${tableName}") + assertEquals(["Rule", "Before", "After"], process.names) + assertTrue(process.rows.size() > 0, "EXPLAIN PLAN PROCESS over Flight returned no rows") + + // 5. Session state set by one request is seen by the next, from either side: a user variable, + // and a session variable read back through a frontend-side SHOW and a backend query. + flight("SET @lifecycle_var = 42") + assertEquals([[42L]], nums(flight("SELECT @lifecycle_var").rows)) + flight("SET query_timeout = 1234") + try { + assertEquals("1234", flight("SHOW VARIABLES LIKE 'query_timeout'").rows[0][1].toString()) + assertEquals([[1234L]], nums(flight("SELECT @@query_timeout").rows)) + } finally { + flight("SET query_timeout = DEFAULT") + } + + // 6. A request may carry several statements; only the last one may return a result. The result + // of a query is on the backend, so the request that ends in a query is fine, and the frontend + // refuses a request that produced a frontend-side result before its last statement. + assertEquals([[2L]], nums(flight("SET @multi = 2; SELECT k FROM ${tableName} WHERE k = 2").rows)) + flightFails("SHOW VARIABLES LIKE 'wait_timeout'; SELECT 1", + "Only be one stmt that returns the result and it is at the end") + + // 7. A failed statement leaves the session usable. + flightFails("SELECT * FROM no_such_table_lifecycle", "does not exist") + assertEquals([[3L]], nums(flight("SELECT count(*) FROM ${tableName}").rows)) + flightFails("SELECT k FROM ${tableName} WHERE", "mismatched input") + assertEquals([[1L, "a"]], nums(flight("SELECT k, v FROM ${tableName} WHERE k = 1").rows)) + + // 8. A DDL over Flight takes effect and answers with the synthesized status row. + sql "DROP TABLE IF EXISTS ${tableName}_ddl" + def ddl = flight("CREATE TABLE ${tableName}_ddl (k INT) DUPLICATE KEY(k) DISTRIBUTED BY HASH(k) BUCKETS 1" + + " PROPERTIES ('replication_num' = '1')") + assertEquals(["StatusResult"], ddl.names) + assertEquals(1, flight("SHOW TABLES LIKE '${tableName}_ddl'").rows.size()) + flight("DROP TABLE ${tableName}_ddl") + assertEquals(0, flight("SHOW TABLES LIKE '${tableName}_ddl'").rows.size()) + + sql "DROP TABLE IF EXISTS ${tableName}" +} diff --git a/regression-test/suites/query_p0/test_multi_statement_response.groovy b/regression-test/suites/query_p0/test_multi_statement_response.groovy new file mode 100644 index 00000000000000..088298fa9d72e5 --- /dev/null +++ b/regression-test/suites/query_p0/test_multi_statement_response.groovy @@ -0,0 +1,439 @@ +// 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 java.net.Socket +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.charset.StandardCharsets +import java.security.MessageDigest +import java.sql.SQLException +import java.sql.Statement + +// What a MySQL client receives for a request that carries several statements. +// +// With CLIENT_MULTI_STATEMENTS negotiated, every statement's response is delivered, the +// intermediate ones flagged SERVER_MORE_RESULTS_EXISTS. Without it Doris still runs every +// statement but delivers only the response of the last one: what an earlier statement wrote is +// dropped when the next statement starts (MysqlProtocolAdapter.beforeStatement), whatever kind of +// statement follows. A request that ends in a SET, an INSERT or a BEGIN therefore answers with that +// statement's OK alone; it used to answer with the packets of the preceding query followed by the +// OK, a stream no client can parse. +// +// Connector/J exercises the second case only: it asks for CLIENT_MULTI_STATEMENTS only when the +// server advertises it, which Doris does not, and it always asks for CLIENT_DEPRECATE_EOF. The +// other combinations are driven through a bare protocol client below, which also sees the response +// packet by packet. +suite("test_multi_statement_response") { + def tableName = "multi_statement_response" + sql "DROP TABLE IF EXISTS ${tableName}" + sql """ + CREATE TABLE ${tableName} (k INT) + DUPLICATE KEY(k) DISTRIBUTED BY HASH(k) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + + // The rows of the current result set of a statement. + def readRows = { Statement statement -> + def rows = [] + statement.getResultSet().withCloseable { resultSet -> + while (resultSet.next()) { + rows.add(resultSet.getObject(1)) + } + } + return rows + } + + // 1. Connector/J, no CLIENT_MULTI_STATEMENTS: only the last statement's response reaches the + // client. The socket timeout turns a response the client cannot parse into a failure + // instead of a hang. + def url = context.config.jdbcUrl + (context.config.jdbcUrl.contains("?") ? "&" : "?") + "socketTimeout=30000" + connect(context.config.jdbcUser, context.config.jdbcPassword, url) { + context.getConnection().createStatement().withCloseable { statement -> + statement.execute("USE ${context.dbName}") + + // The last statement is a query: its result set, and nothing of the query before it. + assertTrue(statement.execute("SELECT 1; SELECT 2")) + assertEquals([2], readRows(statement)) + assertFalse(statement.getMoreResults()) + assertEquals(-1, statement.getUpdateCount()) + + // The last statement is a SET: its OK alone, then the connection is still in step. + assertFalse(statement.execute("SELECT 1; SET @multi_stmt_var = 7")) + assertEquals(0, statement.getUpdateCount()) + assertTrue(statement.execute("SELECT @multi_stmt_var")) + assertEquals([7], readRows(statement)) + + // Two queries before the SET: both dropped. + assertFalse(statement.execute("SELECT 1; SELECT 2; SET @multi_stmt_var = 8")) + assertEquals(0, statement.getUpdateCount()) + assertTrue(statement.execute("SELECT @multi_stmt_var")) + assertEquals([8], readRows(statement)) + + // The last statement is an INSERT: its OK carries the affected rows. + assertFalse(statement.execute("SELECT 1; INSERT INTO ${tableName} VALUES (1)")) + assertEquals(1, statement.getUpdateCount()) + assertTrue(statement.execute("SELECT k FROM ${tableName}")) + assertEquals([1], readRows(statement)) + + // The last statement opens a transaction: its OK alone, and the transaction is open. + assertFalse(statement.execute("SELECT 1; BEGIN")) + assertEquals(0, statement.getUpdateCount()) + statement.execute("ROLLBACK") + + // A SET before the query: the query's result set. + assertTrue(statement.execute("SET @multi_stmt_var = 9; SELECT @multi_stmt_var")) + assertEquals([9], readRows(statement)) + + // A failing statement in the middle ends the request with its error, and the + // connection is still usable afterwards. + try { + statement.execute("SELECT 1; SELECT * FROM no_such_table_multi_stmt; SELECT 2") + throw new AssertionError("a request with a failing statement did not fail") + } catch (SQLException e) { + assertTrue(e.getMessage().contains("does not exist"), e.getMessage()) + } + assertTrue(statement.execute("SELECT 3")) + assertEquals([3], readRows(statement)) + } + } + + // 2. The bare client, every combination of CLIENT_MULTI_STATEMENTS and CLIENT_DEPRECATE_EOF. + String hostPort = context.config.jdbcUrl.substring(context.config.jdbcUrl.indexOf("://") + 3) + hostPort = hostPort.substring(0, hostPort.indexOf("/") >= 0 ? hostPort.indexOf("/") : hostPort.length()) + String host = hostPort.substring(0, hostPort.indexOf(":")) + int port = hostPort.substring(hostPort.indexOf(":") + 1).toInteger() + def rawClient = { boolean multiStatements, boolean deprecateEof -> + return new RawMysqlClient(host, port, context.config.jdbcUser, context.config.jdbcPassword, + context.dbName, multiStatements, deprecateEof) + } + + [true, false].each { deprecateEof -> + // With CLIENT_MULTI_STATEMENTS: every statement's response, in order, each but the last + // flagged SERVER_MORE_RESULTS_EXISTS. + rawClient(true, deprecateEof).withCloseable { client -> + def results = client.query("SELECT 1; SET @multi_stmt_var = 10; SELECT @multi_stmt_var; " + + "INSERT INTO ${tableName} VALUES (2); SELECT k FROM ${tableName} ORDER BY k") + assertEquals(5, results.size(), "deprecate_eof=${deprecateEof}: ${results}") + assertEquals([kind: "rows", rows: [["1"]], more: true], results[0]) + assertEquals([kind: "ok", affectedRows: 0L, more: true], results[1]) + assertEquals([kind: "rows", rows: [["10"]], more: true], results[2]) + assertEquals([kind: "ok", affectedRows: 1L, more: true], results[3]) + assertEquals([kind: "rows", rows: [["1"], ["2"]], more: false], results[4]) + sql "DELETE FROM ${tableName} WHERE k = 2" + + // A failing statement ends the request where it fails: the responses before it were + // delivered, then its error, and the connection is still usable afterwards. + try { + client.query("SELECT 1; SELECT * FROM no_such_table_multi_stmt; SELECT 2") + throw new AssertionError("a request with a failing statement did not fail") + } catch (SQLException e) { + assertTrue(e.getMessage().contains("does not exist"), e.getMessage()) + assertEquals([[kind: "rows", rows: [["1"]], more: true]], client.getPartialResults()) + } + assertEquals([[kind: "rows", rows: [["3"]], more: false]], client.query("SELECT 3")) + } + + // Without CLIENT_MULTI_STATEMENTS: only the last statement's response. + rawClient(false, deprecateEof).withCloseable { client -> + assertEquals([[kind: "rows", rows: [["2"]], more: false]], client.query("SELECT 1; SELECT 2")) + assertEquals([[kind: "ok", affectedRows: 0L, more: false]], + client.query("SELECT 1; SET @multi_stmt_var = 11")) + assertEquals([[kind: "rows", rows: [["11"]], more: false]], client.query("SELECT @multi_stmt_var")) + assertEquals([[kind: "ok", affectedRows: 1L, more: false]], + client.query("SELECT 1; INSERT INTO ${tableName} VALUES (3)")) + assertEquals([[kind: "rows", rows: [["1"], ["3"]], more: false]], + client.query("SELECT k FROM ${tableName} ORDER BY k")) + sql "DELETE FROM ${tableName} WHERE k = 3" + } + } + + sql "DROP TABLE IF EXISTS ${tableName}" +} + +/** + * A bare MySQL protocol client: the handshake with the capability flags a test asks for, COM_QUERY, + * and the response read packet by packet into a list of results, one per statement delivered: + * [kind: "rows", rows: [[column text or null, ...], ...], more: boolean] for a result set and + * [kind: "ok", affectedRows: long, more: boolean] for an OK, where "more" is the + * SERVER_MORE_RESULTS_EXISTS flag of the packet that ended it. An ERR packet is thrown as an + * SQLException; the results delivered before it stay in getPartialResults(). + */ +class RawMysqlClient implements Closeable { + static final int CLIENT_LONG_PASSWORD = 0x00000001 + static final int CLIENT_LONG_FLAG = 0x00000004 + static final int CLIENT_CONNECT_WITH_DB = 0x00000008 + static final int CLIENT_PROTOCOL_41 = 0x00000200 + static final int CLIENT_SECURE_CONNECTION = 0x00008000 + static final int CLIENT_MULTI_STATEMENTS = 0x00010000 + static final int CLIENT_MULTI_RESULTS = 0x00020000 + static final int CLIENT_PLUGIN_AUTH = 0x00080000 + static final int CLIENT_DEPRECATE_EOF = 0x01000000 + static final int SERVER_MORE_RESULTS_EXISTS = 0x0008 + + private final Socket socket + private final DataInputStream input + private final OutputStream output + private final boolean deprecateEof + private int sequenceId = 0 + private List partialResults = [] + + RawMysqlClient(String host, int port, String user, String password, String db, + boolean multiStatements, boolean deprecateEof) { + this.deprecateEof = deprecateEof + socket = new Socket(host, port) + socket.setSoTimeout(30000) + input = new DataInputStream(new BufferedInputStream(socket.getInputStream())) + output = socket.getOutputStream() + + // The greeting: protocol version, server version, connection id, the two parts of the + // 20-byte scramble around the capability flags, charset, status and the plugin name. + ByteBuffer greeting = ByteBuffer.wrap(readPacket()).order(ByteOrder.LITTLE_ENDIAN) + greeting.get() + while (greeting.get() != 0) { + } + greeting.getInt() + byte[] scramble = new byte[20] + greeting.get(scramble, 0, 8) + greeting.get() + greeting.getShort() + greeting.get() + greeting.getShort() + greeting.getShort() + greeting.get() + greeting.position(greeting.position() + 10) + greeting.get(scramble, 8, 12) + + int flags = CLIENT_LONG_PASSWORD | CLIENT_LONG_FLAG | CLIENT_PROTOCOL_41 | CLIENT_SECURE_CONNECTION + | CLIENT_PLUGIN_AUTH | CLIENT_CONNECT_WITH_DB + if (multiStatements) { + flags |= CLIENT_MULTI_STATEMENTS | CLIENT_MULTI_RESULTS + } + if (deprecateEof) { + flags |= CLIENT_DEPRECATE_EOF + } + ByteArrayOutputStream response = new ByteArrayOutputStream() + writeInt4(response, flags) + writeInt4(response, 16777216) + response.write(33) + response.write(new byte[23]) + response.write(user.getBytes(StandardCharsets.UTF_8)) + response.write(0) + byte[] auth = password.isEmpty() ? new byte[0] : nativePasswordScramble(password, scramble) + response.write(auth.length) + response.write(auth) + response.write(db.getBytes(StandardCharsets.UTF_8)) + response.write(0) + response.write("mysql_native_password".getBytes(StandardCharsets.UTF_8)) + response.write(0) + sendPacket(response.toByteArray()) + + byte[] reply = readPacket() + if ((reply[0] & 0xFF) == 0xFE) { + // An auth switch to the same plugin: answer the new scramble. + ByteBuffer request = ByteBuffer.wrap(reply) + request.get() + while (request.get() != 0) { + } + byte[] newScramble = new byte[20] + request.get(newScramble) + sendPacket(password.isEmpty() ? new byte[0] : nativePasswordScramble(password, newScramble)) + reply = readPacket() + } + if ((reply[0] & 0xFF) != 0x00) { + throw new IllegalStateException("handshake failed: " + describeError(reply)) + } + } + + // SHA1(password) XOR SHA1(scramble + SHA1(SHA1(password))) + private static byte[] nativePasswordScramble(String password, byte[] scramble) { + MessageDigest sha1 = MessageDigest.getInstance("SHA-1") + byte[] stage1 = sha1.digest(password.getBytes(StandardCharsets.UTF_8)) + sha1.reset() + byte[] stage2 = sha1.digest(stage1) + sha1.reset() + sha1.update(scramble) + sha1.update(stage2) + byte[] mixed = sha1.digest() + byte[] result = new byte[stage1.length] + for (int i = 0; i < result.length; i++) { + result[i] = (byte) (stage1[i] ^ mixed[i]) + } + return result + } + + List getPartialResults() { + return partialResults + } + + List query(String sql) { + sequenceId = 0 + ByteArrayOutputStream command = new ByteArrayOutputStream() + command.write(0x03) + command.write(sql.getBytes(StandardCharsets.UTF_8)) + sendPacket(command.toByteArray()) + + List results = [] + partialResults = results + while (true) { + byte[] first = readPacket() + int header = first[0] & 0xFF + if (header == 0xFF) { + throw new SQLException(describeError(first)) + } + if (header == 0x00) { + Map ok = parseOk(first) + results.add(ok) + if (!ok.more) { + return results + } + continue + } + int columnCount = (int) readLenenc(ByteBuffer.wrap(first).order(ByteOrder.LITTLE_ENDIAN)) + for (int i = 0; i < columnCount; i++) { + readPacket() + } + if (!deprecateEof) { + byte[] eof = readPacket() + if ((eof[0] & 0xFF) != 0xFE || eof.length != 5) { + throw new IllegalStateException("expected an EOF after the column definitions, got " + + describePacket(eof)) + } + } + List> rows = [] + while (true) { + byte[] packet = readPacket() + int rowHeader = packet[0] & 0xFF + if (rowHeader == 0xFF) { + throw new SQLException(describeError(packet)) + } + if (rowHeader == 0xFE && (deprecateEof ? packet.length >= 7 : packet.length < 9)) { + boolean more = (terminatorStatus(packet) & SERVER_MORE_RESULTS_EXISTS) != 0 + results.add([kind: "rows", rows: rows, more: more]) + if (!more) { + return results + } + break + } + rows.add(parseTextRow(packet, columnCount)) + } + } + } + + private Map parseOk(byte[] packet) { + ByteBuffer buffer = ByteBuffer.wrap(packet).order(ByteOrder.LITTLE_ENDIAN) + buffer.get() + long affectedRows = readLenenc(buffer) + readLenenc(buffer) + int status = buffer.getShort() & 0xFFFF + return [kind: "ok", affectedRows: affectedRows, more: (status & SERVER_MORE_RESULTS_EXISTS) != 0] + } + + // The status flags of the packet that ends a result set: an EOF packet (warnings, status) or, + // when the client deprecated EOF, an OK packet with the 0xFE header. + private int terminatorStatus(byte[] packet) { + ByteBuffer buffer = ByteBuffer.wrap(packet).order(ByteOrder.LITTLE_ENDIAN) + buffer.get() + if (!deprecateEof) { + buffer.getShort() + return buffer.getShort() & 0xFFFF + } + readLenenc(buffer) + readLenenc(buffer) + return buffer.getShort() & 0xFFFF + } + + private static List parseTextRow(byte[] packet, int columnCount) { + ByteBuffer buffer = ByteBuffer.wrap(packet).order(ByteOrder.LITTLE_ENDIAN) + List row = [] + for (int i = 0; i < columnCount; i++) { + if ((buffer.get(buffer.position()) & 0xFF) == 0xFB) { + buffer.get() + row.add(null) + } else { + int length = (int) readLenenc(buffer) + byte[] value = new byte[length] + buffer.get(value) + row.add(new String(value, StandardCharsets.UTF_8)) + } + } + return row + } + + private static long readLenenc(ByteBuffer buffer) { + int first = buffer.get() & 0xFF + if (first < 0xFB) { + return first + } + if (first == 0xFC) { + return buffer.getShort() & 0xFFFFL + } + if (first == 0xFD) { + return (buffer.get() & 0xFFL) | ((buffer.get() & 0xFFL) << 8) | ((buffer.get() & 0xFFL) << 16) + } + return buffer.getLong() + } + + private static String describeError(byte[] packet) { + ByteBuffer buffer = ByteBuffer.wrap(packet).order(ByteOrder.LITTLE_ENDIAN) + buffer.get() + int code = buffer.getShort() & 0xFFFF + byte[] rest = new byte[buffer.remaining()] + buffer.get(rest) + return "error " + code + ": " + new String(rest, StandardCharsets.UTF_8) + } + + private static String describePacket(byte[] packet) { + return packet.length + " bytes starting with 0x" + Integer.toHexString(packet[0] & 0xFF) + } + + private byte[] readPacket() { + byte[] header = new byte[4] + input.readFully(header) + int length = (header[0] & 0xFF) | ((header[1] & 0xFF) << 8) | ((header[2] & 0xFF) << 16) + sequenceId = ((header[3] & 0xFF) + 1) & 0xFF + byte[] payload = new byte[length] + input.readFully(payload) + return payload + } + + private void sendPacket(byte[] payload) { + byte[] header = [(byte) payload.length, (byte) (payload.length >> 8), (byte) (payload.length >> 16), + (byte) sequenceId] + output.write(header) + output.write(payload) + output.flush() + sequenceId = (sequenceId + 1) & 0xFF + } + + private static void writeInt4(ByteArrayOutputStream out, int value) { + out.write(value & 0xFF) + out.write((value >> 8) & 0xFF) + out.write((value >> 16) & 0xFF) + out.write((value >> 24) & 0xFF) + } + + @Override + void close() { + try { + sequenceId = 0 + sendPacket([0x01] as byte[]) + } catch (IOException ignored) { + // the server may have closed first + } + socket.close() + } +} diff --git a/regression-test/suites/query_p0/test_mysql_forward_to_master.groovy b/regression-test/suites/query_p0/test_mysql_forward_to_master.groovy new file mode 100644 index 00000000000000..0c8c0c8e9b81b1 --- /dev/null +++ b/regression-test/suites/query_p0/test_mysql_forward_to_master.groovy @@ -0,0 +1,123 @@ +// 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.doris.regression.suite.ClusterOptions + +// A statement issued over MySQL to a non-master FE is forwarded to the master, and what the client +// receives is the response the master produced for it: the follower tells the master what its +// client negotiated (MysqlProtocolAdapter.fillForwardRequest / restoreFromForwardRequest: the +// capability flags, CLIENT_DEPRECATE_EOF, and for a COM_STMT_EXECUTE the execute packet and whether +// a cursor was asked for), and relays the master's packets as they are +// (MysqlProtocolAdapter.finishCommand -> StmtExecutor.sendProxyQueryResult), reshaped for the +// client's cursor expectations by FEOpExecutor.prepareQueryResultForClient. This needs more than +// one FE, so it can only be exercised in a docker cluster: on a single-FE deployment +// StmtExecutor.shouldForwardToMaster() returns false immediately. +// +// The Arrow Flight SQL side of the same path is test_arrow_flight_forward_to_master. +suite("test_mysql_forward_to_master", "docker") { + def options = new ClusterOptions() + options.setFeNum(2) + options.connectToFollower = true + + docker(options) { + def follower = cluster.getOneFollowerFe() + assertNotNull(follower, "a follower FE is required to exercise the forward-to-master path") + + def tableName = "mysql_forward_to_master" + // `sql` runs on the follower. A DDL is a Redirect command and is forwarded to the master. + sql "DROP TABLE IF EXISTS ${tableName}" + sql """ + CREATE TABLE ${tableName} (k INT) + DUPLICATE KEY(k) DISTRIBUTED BY HASH(k) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + // The same DDL again fails on the master; the master's error is what the client gets. + test { + sql """ + CREATE TABLE ${tableName} (k INT) + DUPLICATE KEY(k) DISTRIBUTED BY HASH(k) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + exception "already exists" + } + sql "INSERT INTO ${tableName} VALUES (1), (2), (3)" + + // 1. A forwarded query answers with the master's result set packets, relayed as they are. + sql "SET force_forward_all_queries = true" + try { + assertEquals([[1], [2], [3]], sql("SELECT k FROM ${tableName} ORDER BY k")) + assertEquals([[1]], sql("SELECT 1")) + assertEquals([[3L]], sql("SELECT count(*) FROM ${tableName}")) + // A query without rows: header, column definitions and the end of the result set only. + assertEquals([], sql("SELECT k FROM ${tableName} WHERE k < 0")) + + // A forwarded SHOW carries a result set of its own (a ShowResultSet, not packets). + def frontends = sql_return_maparray "SHOW FRONTENDS" + assertTrue(frontends.size() >= 2, "expected both FEs, got " + frontends.size()) + assertEquals(1, frontends.count { it.IsMaster == "true" }) + + // A forwarded query that fails on the master: the master's error, and the connection + // is still usable afterwards. + test { + sql "SELECT * FROM no_such_table_forward" + exception "does not exist" + } + assertEquals([[2]], sql("SELECT k FROM ${tableName} WHERE k = 2")) + } finally { + sql "SET force_forward_all_queries = false" + } + + // 2. A forwarded COM_STMT_EXECUTE: the execute packet travels with the request, and the + // result comes back shaped for the cursor the client asked for. A server-side prepared + // statement has to be prepared while forwarding is off (a forwarded COM_STMT_PREPARE + // is refused), so forwarding is switched on between PREPARE and EXECUTE. + String followerUrl = "jdbc:mysql://${follower.host}:${follower.queryPort}/${context.dbName}" + + "?useServerPrepStmts=true&useCursorFetch=true&emulateUnsupportedPstmts=false&socketTimeout=30000" + connect(context.config.jdbcUser, context.config.jdbcPassword, followerUrl) { + def connection = context.getConnection() + connection.createStatement().withCloseable { control -> + control.execute("SET force_forward_all_queries = false") + [0, 1, 10000].each { fetchSize -> + ["SELECT k FROM ${tableName} ORDER BY k".toString(), + "SELECT k FROM ${tableName} WHERE k < 0 ORDER BY k".toString()].each { query -> + connection.prepareStatement(query).withCloseable { prepared -> + assertEquals(com.mysql.cj.jdbc.ServerPreparedStatement, prepared.class) + prepared.setFetchSize(fetchSize) + def readRows = { + def rows = [] + prepared.executeQuery().withCloseable { result -> + while (result.next()) { + rows.add(result.getInt(1)) + } + } + return rows + } + def directRows = readRows() + control.execute("SET force_forward_all_queries = true") + try { + // The same server-prepared statement, executed on the master. + 3.times { assertEquals(directRows, readRows()) } + } finally { + control.execute("SET force_forward_all_queries = false") + } + } + } + } + } + } + } +}