Fix BatchUpdateException when execute INSERT INTO ON DUPLICATE KEY UPDATE in proxy adapter

RELEASE-NOTES.md

@@ -41,6 +41,7 @@
 1. Sharding: Fix avg, sum, min, max function return empty data when no query result return - [#33449](https://github.com/apache/shardingsphere/pull/33449)
 1. Encrypt: Fix merge exception without encrypt rule in database - [#33708](https://github.com/apache/shardingsphere/pull/33708)
 1. SQL Parser: Fix mysql parse zone unreserved keyword error - [#33720](https://github.com/apache/shardingsphere/pull/33720)
+1. Proxy: Fix BatchUpdateException when execute INSERT INTO ON DUPLICATE KEY UPDATE in proxy adapter - [#33796](https://github.com/apache/shardingsphere/pull/33796)
 
 ### Change Logs
 

db-protocol/mysql/src/main/java/org/apache/shardingsphere/db/protocol/mysql/constant/MySQLCapabilityFlag.java

@@ -97,7 +97,7 @@ public static int calculateHandshakeCapabilityFlagsLower() {
      * @return handshake capability flags upper bit
      */
     public static int calculateHandshakeCapabilityFlagsUpper() {
-        return calculateCapabilityFlags(CLIENT_MULTI_STATEMENTS, CLIENT_PLUGIN_AUTH) >> 16;
+        return calculateCapabilityFlags(CLIENT_MULTI_STATEMENTS, CLIENT_PLUGIN_AUTH, CLIENT_MULTI_RESULTS, CLIENT_PS_MULTI_RESULTS) >> 16;
     }
 
     /**

db-protocol/mysql/src/test/java/org/apache/shardingsphere/db/protocol/mysql/constant/MySQLCapabilityFlagTest.java

@@ -36,6 +36,6 @@ void assertCalculateHandshakeCapabilityFlagsLower() {
 
     @Test
     void assertCalculateHandshakeCapabilityFlagsUpper() {
-        assertThat(MySQLCapabilityFlag.calculateHandshakeCapabilityFlagsUpper(), is(0x0009));
+        assertThat(MySQLCapabilityFlag.calculateHandshakeCapabilityFlagsUpper(), is(0x000f));
     }
 }

jdbc/src/main/java/org/apache/shardingsphere/driver/executor/engine/batch/preparedstatement/BatchPreparedStatementExecutor.java

@@ -178,7 +178,7 @@ private int[] accumulate(final List<int[]> executeResults) {
 
     private void accumulate(final int[] executeResult, final int[] addBatchCounts, final JDBCExecutionUnit executionUnit) {
         for (Entry<Integer, Integer> entry : getJDBCAndActualAddBatchCallTimesMap(executionUnit).entrySet()) {
-            int value = null == executeResult ? 0 : executeResult[entry.getValue()];
+            int value = null == executeResult || 0 == executeResult.length ? 0 : executeResult[entry.getValue()];
             addBatchCounts[entry.getKey()] += value;
         }
     }

proxy/backend/core/src/main/java/org/apache/shardingsphere/proxy/backend/response/header/update/MultiStatementsUpdateResponseHeader.java

@@ -0,0 +1,34 @@
+/*
+ * 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.shardingsphere.proxy.backend.response.header.update;
+
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import org.apache.shardingsphere.proxy.backend.response.header.ResponseHeader;
+
+import java.util.Collection;
+
+/**
+ * Multi statements update response header.
+ */
+@RequiredArgsConstructor
+@Getter
+public final class MultiStatementsUpdateResponseHeader implements ResponseHeader {
+
+    private final Collection<UpdateResponseHeader> updateResponseHeaders;
+}

proxy/frontend/type/mysql/src/main/java/org/apache/shardingsphere/proxy/frontend/mysql/command/ServerStatusFlagCalculator.java

@@ -35,9 +35,21 @@ public final class ServerStatusFlagCalculator {
      * @return server status flag
      */
     public static int calculateFor(final ConnectionSession connectionSession) {
+        return calculateFor(connectionSession, true);
+    }
+
+    /**
+     * Calculate server status flag for specified connection.
+     *
+     * @param connectionSession connection session
+     * @param lastPacket last packet
+     * @return server status flag
+     */
+    public static int calculateFor(final ConnectionSession connectionSession, final boolean lastPacket) {
         int result = 0;
         result |= connectionSession.isAutoCommit() ? MySQLStatusFlag.SERVER_STATUS_AUTOCOMMIT.getValue() : 0;
         result |= connectionSession.getTransactionStatus().isInTransaction() ? MySQLStatusFlag.SERVER_STATUS_IN_TRANS.getValue() : 0;
+        result |= lastPacket ? 0 : MySQLStatusFlag.SERVER_MORE_RESULTS_EXISTS.getValue();
         return result;
     }
 }

proxy/frontend/type/mysql/src/main/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/binary/prepare/MySQLComStmtPrepareParameterMarkerExtractor.java

@@ -22,13 +22,15 @@
 import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereColumn;
 import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereSchema;
 import org.apache.shardingsphere.infra.metadata.database.schema.model.ShardingSphereTable;
+import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.ColumnAssignmentSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.InsertValuesSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.ExpressionSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.simple.ParameterMarkerExpressionSegment;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.InsertStatement;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.ListIterator;
@@ -68,10 +70,23 @@ private static List<ShardingSphereColumn> findColumnsOfParameterMarkersForInsert
                 result.add(column);
             }
         }
+        insertStatement.getOnDuplicateKeyColumns().ifPresent(optional -> appendOnDuplicateKeyParameterMarkers(optional.getColumns(), table, result));
         return result;
     }
 
     private static List<String> getColumnNamesOfInsertStatement(final InsertStatement insertStatement, final ShardingSphereTable table) {
         return insertStatement.getColumns().isEmpty() ? table.getColumnNames() : insertStatement.getColumns().stream().map(each -> each.getIdentifier().getValue()).collect(Collectors.toList());
     }
+
+    private static void appendOnDuplicateKeyParameterMarkers(final Collection<ColumnAssignmentSegment> onDuplicateKeyColumns,
+                                                             final ShardingSphereTable table, final List<ShardingSphereColumn> result) {
+        for (ColumnAssignmentSegment each : onDuplicateKeyColumns) {
+            if (!(each.getValue() instanceof ParameterMarkerExpressionSegment)) {
+                continue;
+            }
+            String columnName = each.getColumns().iterator().next().getIdentifier().getValue();
+            ShardingSphereColumn column = table.getColumn(columnName);
+            result.add(column);
+        }
+    }
 }

proxy/frontend/type/mysql/src/main/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLComQueryPacketExecutor.java

@@ -31,6 +31,7 @@
 import org.apache.shardingsphere.proxy.backend.handler.ProxySQLComQueryParser;
 import org.apache.shardingsphere.proxy.backend.response.header.ResponseHeader;
 import org.apache.shardingsphere.proxy.backend.response.header.query.QueryResponseHeader;
+import org.apache.shardingsphere.proxy.backend.response.header.update.MultiStatementsUpdateResponseHeader;
 import org.apache.shardingsphere.proxy.backend.response.header.update.UpdateResponseHeader;
 import org.apache.shardingsphere.proxy.backend.session.ConnectionSession;
 import org.apache.shardingsphere.proxy.frontend.command.executor.QueryCommandExecutor;
@@ -39,10 +40,12 @@
 import org.apache.shardingsphere.proxy.frontend.mysql.command.query.builder.ResponsePacketBuilder;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.DeleteStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.InsertStatement;
 import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.UpdateStatement;
 
 import java.sql.SQLException;
 import java.util.Collection;
+import java.util.LinkedList;
 
 /**
  * COM_QUERY command packet executor for MySQL.
@@ -78,7 +81,11 @@ private boolean isMultiStatementsEnabled(final ConnectionSession connectionSessi
     }
 
     private boolean isSuitableMultiStatementsSQLStatement(final SQLStatement sqlStatement) {
-        return sqlStatement instanceof UpdateStatement || sqlStatement instanceof DeleteStatement;
+        return containsInsertOnDuplicateKey(sqlStatement) || sqlStatement instanceof UpdateStatement || sqlStatement instanceof DeleteStatement;
+    }
+
+    private boolean containsInsertOnDuplicateKey(final SQLStatement sqlStatement) {
+        return sqlStatement instanceof InsertStatement && ((InsertStatement) sqlStatement).getOnDuplicateKeyColumns().isPresent();
     }
 
     @Override
@@ -88,6 +95,9 @@ public Collection<DatabasePacket> execute() throws SQLException {
             return processQuery((QueryResponseHeader) responseHeader);
         }
         responseType = ResponseType.UPDATE;
+        if (responseHeader instanceof MultiStatementsUpdateResponseHeader) {
+            return processMultiStatementsUpdate((MultiStatementsUpdateResponseHeader) responseHeader);
+        }
         return processUpdate((UpdateResponseHeader) responseHeader);
     }
 
@@ -100,6 +110,16 @@ private Collection<DatabasePacket> processUpdate(final UpdateResponseHeader upda
         return ResponsePacketBuilder.buildUpdateResponsePackets(updateResponseHeader, ServerStatusFlagCalculator.calculateFor(connectionSession));
     }
 
+    private Collection<DatabasePacket> processMultiStatementsUpdate(final MultiStatementsUpdateResponseHeader responseHeader) {
+        Collection<DatabasePacket> result = new LinkedList<>();
+        int index = 0;
+        for (UpdateResponseHeader each : responseHeader.getUpdateResponseHeaders()) {
+            boolean lastPacket = ++index == responseHeader.getUpdateResponseHeaders().size();
+            result.addAll(ResponsePacketBuilder.buildUpdateResponsePackets(each, ServerStatusFlagCalculator.calculateFor(connectionSession, lastPacket)));
+        }
+        return result;
+    }
+
     @Override
     public boolean next() throws SQLException {
         return proxyBackendHandler.next();