(shardingsphere) branch master updated: Fix multi statements with specified database name (#29679)

[duanzhengqiang]

This is an automated email from the ASF dual-hosted git repository.

duanzhengqiang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git


The following commit(s) were added to refs/heads/master by this push:
     new e8a8666bb29 Fix multi statements with specified database name (#29679)
e8a8666bb29 is described below

commit e8a8666bb293786cade973ca988d28a1a3cd6e23
Author: ZhangCheng <chengzh...@apache.org>
AuthorDate: Tue Jan 9 09:21:39 2024 +0800

    Fix multi statements with specified database name (#29679)
---
 .../query/text/query/MySQLMultiStatementsHandler.java     |  6 +++++-
 .../query/text/query/MySQLMultiStatementsHandlerTest.java | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git 
a/proxy/frontend/type/mysql/src/main/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLMultiStatementsHandler.java
 
b/proxy/frontend/type/mysql/src/main/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLMultiStatementsHandler.java
index 1b436dcc70b..f9b28e0ddd6 100644
--- 
a/proxy/frontend/type/mysql/src/main/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLMultiStatementsHandler.java
+++ 
b/proxy/frontend/type/mysql/src/main/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLMultiStatementsHandler.java
@@ -101,7 +101,11 @@ public final class MySQLMultiStatementsHandler implements 
ProxyBackendHandler {
         SQLParserEngine sqlParserEngine = getSQLParserEngine();
         for (String each : extractMultiStatements(pattern, sql)) {
             SQLStatement eachSQLStatement = sqlParserEngine.parse(each, false);
-            ExecutionContext executionContext = 
createExecutionContext(createQueryContext(each, eachSQLStatement));
+            QueryContext queryContext = createQueryContext(each, 
eachSQLStatement);
+            if (null == connectionSession.getQueryContext()) {
+                connectionSession.setQueryContext(queryContext);
+            }
+            ExecutionContext executionContext = 
createExecutionContext(queryContext);
             if (null == anyExecutionContext) {
                 anyExecutionContext = executionContext;
             }
diff --git 
a/proxy/frontend/type/mysql/src/test/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLMultiStatementsHandlerTest.java
 
b/proxy/frontend/type/mysql/src/test/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLMultiStatementsHandlerTest.java
index e75361fd28b..109c794622d 100644
--- 
a/proxy/frontend/type/mysql/src/test/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLMultiStatementsHandlerTest.java
+++ 
b/proxy/frontend/type/mysql/src/test/java/org/apache/shardingsphere/proxy/frontend/mysql/command/query/text/query/MySQLMultiStatementsHandlerTest.java
@@ -85,6 +85,21 @@ class MySQLMultiStatementsHandlerTest {
         assertThat(actualHeader.getSqlStatement(), is(expectedStatement));
     }
     
+    @Test
+    void assertExecuteWithSpecifiedDatabaseName() throws SQLException {
+        String sql = "update foo_db.t set v=v+1 where id=1;update foo_db.t set 
v=v+1 where id=2;update foo_db.t set v=v+1 where id=3";
+        ConnectionSession connectionSession = mockConnectionSession();
+        MySQLUpdateStatement expectedStatement = 
mock(MySQLUpdateStatement.class);
+        ContextManager contextManager = mockContextManager();
+        
when(ProxyContext.getInstance().getContextManager()).thenReturn(contextManager);
+        ResponseHeader actual = new 
MySQLMultiStatementsHandler(connectionSession, expectedStatement, 
sql).execute();
+        assertThat(actual, instanceOf(UpdateResponseHeader.class));
+        UpdateResponseHeader actualHeader = (UpdateResponseHeader) actual;
+        assertThat(actualHeader.getUpdateCount(), is(3L));
+        assertThat(actualHeader.getLastInsertId(), is(0L));
+        assertThat(actualHeader.getSqlStatement(), is(expectedStatement));
+    }
+    
     private ConnectionSession mockConnectionSession() throws SQLException {
         ConnectionSession result = mock(ConnectionSession.class, 
RETURNS_DEEP_STUBS);
         when(result.getDatabaseName()).thenReturn("foo_db");