MysqlProtocolAdapter.java

// 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.mysql.protocol;

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;
import org.apache.doris.mysql.MysqlResultSetEndPacket;
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;
import org.apache.doris.qe.ConnectContext.ConnectType;
import org.apache.doris.qe.ConnectPoolMgr;
import org.apache.doris.qe.ConnectProcessor;
import org.apache.doris.qe.ConnectScheduler;
import org.apache.doris.qe.QueryState.MysqlStateType;
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;
import org.apache.logging.log4j.Logger;

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
 * with it, the handshake and SSL state, and the COM_STMT_EXECUTE packet being processed.
 *
 * <p>The channel decides what kind of MySQL connection this is: a {@link MysqlChannel} over a live
 * socket for a client, a {@code ProxyMysqlChannel} for the context the master builds to replay a
 * forwarded statement, or a {@code DummyMysqlChannel} for an internal context that never talks to
 * a client.
 */
public class MysqlProtocolAdapter implements ProtocolAdapter {
    private static final Logger LOG = LogManager.getLogger(MysqlProtocolAdapter.class);
    private static final String SSL_PROTOCOL = "TLS";

    private final MysqlChannel channel;
    // the protocol capability which server say it can support
    private final MysqlCapability serverCapability = MysqlCapability.DEFAULT_CAPABILITY;
    // the protocol capability after server and client negotiate
    private volatile MysqlCapability capability;
    // This context is used for SSL connection between server and mysql client.
    private final MysqlSslContext sslContext = new MysqlSslContext(SSL_PROTOCOL);
    private MysqlHandshakePacket handshakePacket;
    // The COM_STMT_EXECUTE packet being processed, forwarded as-is when the statement goes to the master.
    private ByteBuffer prepareExecuteBuffer;
    // Whether the current COM_STMT_EXECUTE requested a server-side read-only cursor.
    private boolean cursorFetchRequested;

    public MysqlProtocolAdapter(MysqlChannel channel) {
        this.channel = channel;
    }

    /** The MySQL side of {@code ctx}; throws if the connection speaks another protocol. */
    public static MysqlProtocolAdapter of(ConnectContext ctx) {
        ProtocolAdapter adapter = ctx.getProtocolAdapter();
        if (adapter instanceof MysqlProtocolAdapter) {
            return (MysqlProtocolAdapter) adapter;
        }
        throw new IllegalStateException("not a MySQL connection: "
                + (adapter == null ? "no protocol adapter" : adapter.type()));
    }

    @Override
    public ConnectType type() {
        return ConnectType.MYSQL;
    }

    @Override
    public String remoteHostPortString(ConnectContext ctx) {
        return channel.getRemoteHostPortString();
    }

    @Override
    public TResultSinkType resultSinkType() {
        return TResultSinkType.MYSQL_PROTOCOL;
    }

    @Override
    public MysqlResultSender resultSender(ConnectContext ctx) {
        return new MysqlResultSender(ctx, this);
    }

    @Override
    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 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<ByteBuffer> proxyResultPackets() {
        return ((ProxyMysqlChannel) channel).getProxyResultBufferList();
    }

    @Override
    public ConnectPoolMgr connectPool(ConnectScheduler scheduler) {
        return scheduler.getConnectPoolMgr();
    }

    /**
     * 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 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)
            throws IOException {
        if (stmtIndex != stmtCount - 1) {
            ctx.getState().serverStatus |= MysqlServerStatusFlag.SERVER_MORE_RESULTS_EXISTS;
            if (ctx.getState().getStateType() != MysqlStateType.ERR && channel.clientMultiStatements()) {
                finishCommand(ctx, executor);
            }
        }
        return true;
    }

    /**
     * Sends the response of the command: the OK, EOF or ERR packet that {@code ctx.getState()}
     * describes, or, for a statement that was forwarded to the master, the packets the master
     * produced. {@code executor} is null for a command that ran no statement.
     */
    public void finishCommand(ConnectContext ctx, StmtExecutor executor) throws IOException {
        LOG.debug("Finalize command for query {}", DebugUtil.printId(ctx.queryId()));
        ByteBuffer packet;
        if (executor != null && executor.hasForwardedToMaster()
                && ctx.getState().getStateType() != MysqlStateType.ERR) {
            ShowResultSet resultSet = executor.getShowResultSet();
            if (resultSet == null) {
                executor.sendProxyQueryResult();
                packet = executor.getOutputPacket();
            } else {
                executor.sendResultSet(resultSet);
                packet = responsePacket(ctx);
            }
        } else {
            packet = responsePacket(ctx);
        }

        if (packet == null) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("packet == null");
            }
            return;
        }

        LOG.debug("Send to mysql channel for query {}", DebugUtil.printId(ctx.queryId()));
        channel.sendAndFlush(packet);
        // note(wb) we should write profile after return result to mysql client
        // because write profile maybe take too much time
        // explain query stmt do not have profile
        if (executor != null && executor.getParsedStmt() != null && !executor.getParsedStmt().isExplain()
                && (executor.getParsedStmt() instanceof LogicalPlanAdapter)) {
            executor.updateProfile(true);
            StatsErrorEstimator statsErrorEstimator = ctx.getStatsErrorEstimator();
            if (statsErrorEstimator != null) {
                statsErrorEstimator.updateProfile(ctx.queryId());
            }
        }
        LOG.debug("End finalizing command for query {}", DebugUtil.printId(ctx.queryId()));
    }

    /**
     * The packet that answers the command according to {@code ctx.getState()}: null when the
     * command needs no response or the handler already sent one.
     */
    public ByteBuffer responsePacket(ConnectContext ctx) {
        MysqlPacket packet;
        // When CLIENT_DEPRECATE_EOF is set and the state is EOF (end of result set),
        // we need to send a "ResultSet OK" packet (0xFE header with payload > 5 bytes)
        // instead of the traditional EOF packet. This is required by the MySQL protocol
        // and expected by MySQL Connector/J 9.5.0+.
        if (ctx.getState().getStateType() == MysqlStateType.EOF && channel.clientDeprecatedEOF()) {
            packet = new MysqlResultSetEndPacket(ctx.getState());
        } else {
            packet = ctx.getState().toResponsePacket();
        }
        if (packet == null) {
            // possible two cases:
            // 1. handler has send request
            // 2. this command need not to send response
            return null;
        }

        MysqlSerializer serializer = channel.getSerializer();
        serializer.reset();
        packet.writeTo(serializer);
        return serializer.toByteBuffer();
    }

    @Override
    public void afterStatement(ConnectContext ctx) {
        cursorFetchRequested = false;
    }

    @Override
    public void closeConnection(ConnectContext ctx) {
        channel.close();
    }

    public MysqlChannel getChannel() {
        return channel;
    }

    public MysqlCapability getServerCapability() {
        return serverCapability;
    }

    public MysqlCapability getCapability() {
        return capability;
    }

    public void setCapability(MysqlCapability capability) {
        this.capability = capability;
    }

    public MysqlSslContext getSslContext() {
        return sslContext;
    }

    public void setHandshakePacket(MysqlHandshakePacket handshakePacket) {
        this.handshakePacket = handshakePacket;
    }

    public byte[] getAuthPluginData() {
        return handshakePacket == null ? null : handshakePacket.getAuthPluginData();
    }

    public ByteBuffer getPrepareExecuteBuffer() {
        return prepareExecuteBuffer;
    }

    public void setPrepareExecuteBuffer(ByteBuffer prepareExecuteBuffer) {
        this.prepareExecuteBuffer = prepareExecuteBuffer;
    }

    public boolean isCursorFetchRequested() {
        return cursorFetchRequested;
    }

    public void setCursorFetchRequested(boolean cursorFetchRequested) {
        this.cursorFetchRequested = cursorFetchRequested;
    }

    /**
     * Whether the client will consume the first OK packet after the column definitions of the
     * result it asked a cursor for. Connector/J before 9.5 does, while probing whether the cursor
     * was created; Doris never creates one, so the result must carry one more end marker or an
     * empty result would leave the client waiting.
     */
    public boolean clientConsumesCursorMetadataTerminator(ConnectContext ctx) {
        return cursorFetchRequested
                && MysqlCursorFetchCompatibility.resolve(ctx.getConnectAttributes())
                        != MysqlCursorFetchCompatibility.Behavior.STANDARD;
    }

    // The xnio read loop. A command is read and processed while reads are suspended, which is what
    // serializes the commands of one connection.
    public void startAcceptQuery(ConnectContext ctx, ConnectProcessor connectProcessor) {
        channel.startAcceptQuery(ctx, connectProcessor);
    }

    public void suspendAcceptQuery() {
        channel.suspendAcceptQuery();
    }

    public void resumeAcceptQuery() {
        channel.resumeAcceptQuery();
    }

    public void stopAcceptQuery() throws IOException {
        channel.stopAcceptQuery();
    }
}