#3556, carry PR#5194 to dev-4.x (#5733)

apache · May 21, 2020 · 308f09f · 308f09f
1 parent 8ad915c
commit 308f09f
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 10 deletions.
diff --git a/...he/shardingsphere/shardingproxy/backend/communication/jdbc/execute/JDBCExecuteEngine.java b/...he/shardingsphere/shardingproxy/backend/communication/jdbc/execute/JDBCExecuteEngine.java
@@ -18,7 +18,6 @@
 package org.apache.shardingsphere.shardingproxy.backend.communication.jdbc.execute;
 
 import lombok.Getter;
-import org.apache.shardingsphere.underlying.executor.context.ExecutionContext;
 import org.apache.shardingsphere.sharding.execute.sql.StatementExecuteUnit;
 import org.apache.shardingsphere.sharding.execute.sql.execute.SQLExecuteTemplate;
 import org.apache.shardingsphere.sharding.execute.sql.execute.threadlocal.ExecutorExceptionHandler;
@@ -36,8 +35,11 @@
 import org.apache.shardingsphere.shardingproxy.backend.response.update.UpdateResponse;
 import org.apache.shardingsphere.shardingproxy.context.ShardingProxyContext;
 import org.apache.shardingsphere.sql.parser.binder.statement.SQLStatementContext;
+import org.apache.shardingsphere.sql.parser.sql.statement.dml.DeleteStatement;
 import org.apache.shardingsphere.sql.parser.sql.statement.dml.InsertStatement;
+import org.apache.shardingsphere.sql.parser.sql.statement.dml.UpdateStatement;
 import org.apache.shardingsphere.underlying.common.config.properties.ConfigurationPropertyKey;
+import org.apache.shardingsphere.underlying.executor.context.ExecutionContext;
 import org.apache.shardingsphere.underlying.executor.engine.InputGroup;
 
 import java.sql.SQLException;
@@ -75,12 +77,23 @@ public BackendResponse execute(final ExecutionContext executionContext) throws S
         boolean isExceptionThrown = ExecutorExceptionHandler.isExceptionThrown();
         Collection<InputGroup<StatementExecuteUnit>> inputGroups = sqlExecutePrepareTemplate.getExecuteUnitGroups(
                 executionContext.getExecutionUnits(), new ProxyJDBCExecutePrepareCallback(backendConnection, jdbcExecutorWrapper, isReturnGeneratedKeys));
-        Collection<ExecuteResponse> executeResponses = sqlExecuteTemplate.execute((Collection) inputGroups, 
+        Collection<ExecuteResponse> executeResponses = sqlExecuteTemplate.execute((Collection) inputGroups,
                 new ProxySQLExecuteCallback(sqlStatementContext, backendConnection, jdbcExecutorWrapper, isExceptionThrown, isReturnGeneratedKeys, true),
                 new ProxySQLExecuteCallback(sqlStatementContext, backendConnection, jdbcExecutorWrapper, isExceptionThrown, isReturnGeneratedKeys, false));
         ExecuteResponse executeResponse = executeResponses.iterator().next();
-        return executeResponse instanceof ExecuteQueryResponse
-                ? getExecuteQueryResponse(((ExecuteQueryResponse) executeResponse).getQueryHeaders(), executeResponses) : new UpdateResponse(executeResponses);
+        if (executeResponse instanceof ExecuteQueryResponse) {
+            return getExecuteQueryResponse(((ExecuteQueryResponse) executeResponse).getQueryHeaders(), executeResponses);
+        } else {
+            UpdateResponse updateResponse = new UpdateResponse(executeResponses);
+            if (sqlStatementContext.getSqlStatement() instanceof InsertStatement) {
+                updateResponse.setType("INSERT");
+            } else if (sqlStatementContext.getSqlStatement() instanceof DeleteStatement) {
+                updateResponse.setType("DELETE");
+            } else if (sqlStatementContext.getSqlStatement() instanceof UpdateStatement) {
+                updateResponse.setType("UPDATE");
+            }
+            return updateResponse;
+        }
     }
 
     private BackendResponse getExecuteQueryResponse(final List<QueryHeader> queryHeaders, final Collection<ExecuteResponse> executeResponses) {

diff --git a/.../java/org/apache/shardingsphere/shardingproxy/backend/response/update/UpdateResponse.java b/.../java/org/apache/shardingsphere/shardingproxy/backend/response/update/UpdateResponse.java
@@ -18,6 +18,7 @@
 package org.apache.shardingsphere.shardingproxy.backend.response.update;
 
 import lombok.Getter;
+import lombok.Setter;
 import org.apache.shardingsphere.shardingproxy.backend.communication.jdbc.execute.response.ExecuteResponse;
 import org.apache.shardingsphere.shardingproxy.backend.communication.jdbc.execute.response.ExecuteUpdateResponse;
 import org.apache.shardingsphere.shardingproxy.backend.response.BackendResponse;
@@ -40,6 +41,10 @@ public final class UpdateResponse implements BackendResponse {
     @Getter
     private long updateCount;
 
+    @Getter
+    @Setter
+    private String type;
+
     public UpdateResponse() {
         this(Collections.emptyList());
     }

diff --git a/...hardingproxy/frontend/postgresql/command/query/binary/bind/PostgreSQLComBindExecutor.java b/...hardingproxy/frontend/postgresql/command/query/binary/bind/PostgreSQLComBindExecutor.java
@@ -103,7 +103,7 @@ private PostgreSQLErrorResponsePacket createErrorPacket(final ErrorResponse erro
     }
 
     private PostgreSQLCommandCompletePacket createUpdatePacket(final UpdateResponse updateResponse) {
-        return new PostgreSQLCommandCompletePacket();
+        return new PostgreSQLCommandCompletePacket(updateResponse.getType(), updateResponse.getUpdateCount());
     }
 
     private Optional<PostgreSQLRowDescriptionPacket> createQueryPacket(final QueryResponse queryResponse) {

diff --git a/...here/shardingproxy/frontend/postgresql/command/query/text/PostgreSQLComQueryExecutor.java b/...here/shardingproxy/frontend/postgresql/command/query/text/PostgreSQLComQueryExecutor.java
@@ -88,7 +88,7 @@ private PostgreSQLErrorResponsePacket createErrorPacket(final ErrorResponse erro
     }
 
     private PostgreSQLCommandCompletePacket createUpdatePacket(final UpdateResponse updateResponse) {
-        return new PostgreSQLCommandCompletePacket();
+        return new PostgreSQLCommandCompletePacket(updateResponse.getType(), updateResponse.getUpdateCount());
     }
 
     private Optional<PostgreSQLRowDescriptionPacket> createQueryPacket(final QueryResponse queryResponse) {

diff --git a/...ngsphere/database/protocol/postgresql/packet/generic/PostgreSQLCommandCompletePacket.java b/...ngsphere/database/protocol/postgresql/packet/generic/PostgreSQLCommandCompletePacket.java
@@ -30,13 +30,22 @@ public final class PostgreSQLCommandCompletePacket implements PostgreSQLPacket {
     @Getter
     private final char messageType = PostgreSQLCommandPacketType.COMMAND_COMPLETE.getValue();
 
-    private final String sqlCommand = "";
+    private final String sqlCommand;
 
-    private final int rowCount = 0;
+    private final long rowCount;
+
+    public PostgreSQLCommandCompletePacket() {
+        sqlCommand = "";
+        rowCount = 0;
+    }
+
+    public PostgreSQLCommandCompletePacket(final String sqlCommand, final long rowCount) {
+        this.sqlCommand = sqlCommand;
+        this.rowCount = rowCount;
+    }
 
     @Override
     public void write(final PostgreSQLPacketPayload payload) {
-        // TODO payload.writeStringNul(sqlCommand + " " + rowCount);
-        payload.writeStringNul("");
+        payload.writeStringNul(sqlCommand + " " + rowCount);
     }
 }