branch-2.1: [fix](suites) Fix atomic restore alter suite with master_sql #46550 (#46653)

Cherry-picked from #46550

Co-authored-by: walter <maochuan@selectdb.com>

regression-test/framework/src/main/groovy/org/apache/doris/regression/suite/Suite.groovy

@@ -359,6 +359,10 @@ class Suite implements GroovyInterceptable {
         }
     }
 
+    List<List<Object>> master_sql(String sqlStr, boolean isOrder = false) {
+        return sql_impl(context.getMasterConnection(), sqlStr, isOrder)
+    }
+
     List<List<Object>> multi_sql(String sqlStr, boolean isOrder = false) {
         String[] sqls = sqlStr.split(";")
         def result = new ArrayList<Object>();
@@ -448,6 +452,10 @@ class Suite implements GroovyInterceptable {
         return sql_return_maparray_impl(sqlStr, context.getConnection())
     }
 
+    def master_sql_return_maparray(String sqlStr) {
+        return sql_return_maparray_impl(sqlStr, context.getMasterConnection())
+    }
+
     def arrow_flight_sql_return_maparray(String sqlStr) {
         return sql_return_maparray_impl((String) ("USE ${context.dbName};" + sqlStr), context.getArrowFlightSqlConnection())
     }
@@ -1066,6 +1074,8 @@ class Suite implements GroovyInterceptable {
                     tupleResult = JdbcUtils.executeToStringList(context.getArrowFlightSqlConnection(), (PreparedStatement) arg)
                 } else if (tag.contains("target_sql")) {
                     tupleResult = JdbcUtils.executeToStringList(context.getTargetConnection(this), (PreparedStatement) arg)
+                } else if (tag.contains("master_sql")) {
+                    tupleResult = JdbcUtils.executeToStringList(context.getMasterConnection(), (PreparedStatement) arg)
                 } else {
                     tupleResult = JdbcUtils.executeToStringList(context.getConnection(),  (PreparedStatement) arg)
                 }
@@ -1079,6 +1089,8 @@ class Suite implements GroovyInterceptable {
                             (String) ("USE ${context.dbName};" + (String) arg))
                 } else if (tag.contains("target_sql")) {
                     tupleResult = JdbcUtils.executeToStringList(context.getTargetConnection(this), (String) arg)
+                } else if (tag.contains("master_sql")) {
+                    tupleResult = JdbcUtils.executeToStringList(context.getMasterConnection(), (PreparedStatement) arg)
                 } else {
                     tupleResult = JdbcUtils.executeToStringList(context.getConnection(),  (String) arg)
                 }

regression-test/framework/src/main/groovy/org/apache/doris/regression/suite/SuiteContext.groovy

@@ -44,6 +44,7 @@ class SuiteContext implements Closeable {
     public final String group
     public final String dbName
     public final ThreadLocal<ConnectionInfo> threadLocalConn = new ThreadLocal<>()
+    public final ThreadLocal<ConnectionInfo> threadLocalMasterConn = new ThreadLocal<>()
     public final ThreadLocal<ConnectionInfo> threadArrowFlightSqlConn = new ThreadLocal<>()
     public final ThreadLocal<Connection> threadHive2DockerConn = new ThreadLocal<>()
     public final ThreadLocal<Connection> threadHive3DockerConn = new ThreadLocal<>()
@@ -145,13 +146,26 @@ class SuiteContext implements Closeable {
         if (threadConnInfo == null) {
             threadConnInfo = new ConnectionInfo()
             threadConnInfo.conn = config.getConnectionByDbName(dbName)
-            threadConnInfo.username = config.jdbcUser 
+            threadConnInfo.username = config.jdbcUser
             threadConnInfo.password = config.jdbcPassword
             threadLocalConn.set(threadConnInfo)
         }
         return threadConnInfo.conn
     }
 
+    // like getConnection, but connect to FE master
+    Connection getMasterConnection() {
+        def threadConnInfo = threadLocalMasterConn.get()
+        if (threadConnInfo == null) {
+            threadConnInfo = new ConnectionInfo()
+            threadConnInfo.conn = getMasterConnectionByDbName(dbName)
+            threadConnInfo.username = config.jdbcUser
+            threadConnInfo.password = config.jdbcPassword
+            threadLocalMasterConn.set(threadConnInfo)
+        }
+        return threadConnInfo.conn
+    }
+
     Connection getArrowFlightSqlConnection() {
         def threadConnInfo = threadArrowFlightSqlConn.get()
         if (threadConnInfo == null) {
@@ -316,6 +330,27 @@ class SuiteContext implements Closeable {
         }
     }
 
+    Connection getMasterConnectionByDbName(String dbName) {
+        def result = JdbcUtils.executeToMapArray(getConnection(), "SHOW FRONTENDS")
+        def master = null
+        for (def row : result) {
+            if (row.IsMaster == "true") {
+                master = row
+                break
+            }
+        }
+        if (master) {
+            log.info("master found: ${master.Host}:${master.HttpPort}")
+            def url = Config.buildUrlWithDb(master.Host as String, master.QueryPort as Integer, dbName)
+            def username = config.jdbcUser
+            def password = config.jdbcPassword
+
+            return DriverManager.getConnection(url, username, password)
+        } else {
+            throw new Exception("No master found to reconnect")
+        }
+    }
+
     def reconnectToMasterFe = { ->
         log.info("Reconnecting to a new master frontend...")
         def result = JdbcUtils.executeToMapArray(getConnection(), "SHOW FRONTENDS")
@@ -468,6 +503,16 @@ class SuiteContext implements Closeable {
             }
         }
 
+        ConnectionInfo master_conn = threadLocalMasterConn.get()
+        if (master_conn != null) {
+            threadLocalMasterConn.remove()
+            try {
+                master_conn.conn.close()
+            } catch (Throwable t) {
+                log.warn("Close master connection failed", t)
+            }
+        }
+
         ConnectionInfo arrow_flight_sql_conn = threadArrowFlightSqlConn.get()
         if (arrow_flight_sql_conn != null) {
             threadArrowFlightSqlConn.remove()

regression-test/suites/backup_restore/test_backup_restore_atomic_with_alter.groovy

@@ -125,7 +125,7 @@ suite("test_backup_restore_atomic_with_alter", "backup_restore") {
     sql "SYNC"
 
     // 0. table_1 has in_atomic_restore property
-    def show_result = sql """ SHOW CREATE TABLE ${dbName}.${tableNamePrefix}_1 """
+    def show_result = master_sql """ SHOW CREATE TABLE ${dbName}.${tableNamePrefix}_1 """
     logger.info("SHOW CREATE TABLE ${tableNamePrefix}_1: ${show_result}")
     assertTrue(show_result[0][1].contains("in_atomic_restore"))
 
@@ -230,7 +230,7 @@ suite("test_backup_restore_atomic_with_alter", "backup_restore") {
     sql "SYNC"
 
     // 5. The restore job is cancelled, the in_atomic_restore property has been removed.
-    show_result = sql """ SHOW CREATE TABLE ${dbName}.${tableNamePrefix}_1 """
+    show_result = master_sql """ SHOW CREATE TABLE ${dbName}.${tableNamePrefix}_1 """
     logger.info("SHOW CREATE TABLE ${tableNamePrefix}_1: ${show_result}")
     assertFalse(show_result[0][1].contains("in_atomic_restore"))