test: support captured runtime values in SQL cases

README.md

@@ -105,6 +105,7 @@ Sometimes, to achieve some specific purposes, such as pausing or creating a new
 | -- @skip:issue#{IssueNo.}                                     | If set, the whole script file will be skipped, and not be executed any more for issue{IssueNo.}                                                                                                       |
 | -- @bvt:issue#{IssueNo.}<br/>-- @bvt:issue                    | The sql statements between those two tags will be not executed for issue{IssueNo.}                                                                                                                    |
 | -- @sleep:{time}                                              | The mo-tester will wait for {time} s                                                                                                                                                                  |
+| -- @capture:{name}                                            | Execute the following scalar query, store its first non-null value as `{name}`, and replace later `{{name}}` placeholders before execution. The capture query itself is not result-compared. |
 | -- @wait_expect({interval}, {timeout})                        | Intelligent polling mechanism that repeatedly executes the SQL until the result matches the expected output or timeout is reached. `{interval}` is the polling interval in seconds (how often to retry), and `{timeout}` is the maximum wait time in seconds. This replaces fixed sleep times with smart waiting.<br/><br/>**Example:**<br/>`-- @wait_expect(1, 10)`<br/>`SELECT * FROM orders WHERE status = 'completed';`<br/><br/>This will poll every 1 second for up to 10 seconds until the query returns the expected result. Useful for testing asynchronous operations, data replication, or eventual consistency scenarios.<br/><br/>**Modes:**<br/>- **run mode**: Polls until result matches or timeout<br/>- **genrs mode**: Sleeps for timeout seconds, then executes once<br/><br/>**Features:**<br/>- Supports query results (SELECT)<br/>- Supports error matching (waiting for errors to appear/disappear)<br/>- Detailed logging of retry attempts and timing |
 | -- @session:id=2&user=root&password=111<br/> -- @session      | The mo-tester will create a new connetion to execute sql statements between those two tags.<br/>Default value of id is 1, max is 10.<br/>Defualt value of user and password is configured in `mo.yml`. |
 | -- @sortkey:                                                  | If the result is sorted, need set this tag for the sql statement. e.g.<br/> -- @sortkey:0,1: means sort keys are first column and second colum.                                                       |

src/main/java/io/mo/cases/SqlCommand.java

@@ -124,6 +124,11 @@ public class SqlCommand {
     @Setter
     private int waitExpectTimeout = 20; // wait_expect 超时时间（秒）
 
+    /** Name used by @capture to expose this query's first scalar value. */
+    @Getter
+    @Setter
+    private String captureName = null;
+
     public SqlCommand() {
         command = new StringBuffer();
         testResult = new TestResult();

src/main/java/io/mo/constant/COMMON.java

@@ -28,6 +28,7 @@ public class COMMON {
     public static String COLUMN_SEPARATOR_FLAG = "-- @separator:";
     public static String WAIT_FLAG = "-- @wait:";
     public static String WAIT_EXPECT_FLAG = "-- @wait_expect";
+    public static String CAPTURE_FLAG = "-- @capture:";
 
     // if result is type of error, and not unique, can use this flag to regular
     // match
@@ -90,4 +91,4 @@ public class COMMON {
 
     public static double SCALE_TOLERABLE_ERROR = 0.0000009;
     public static double INT_TOLERABLE_ERROR = 0.000000000000001;
-}
+}

src/main/java/io/mo/db/Executor.java

@@ -18,8 +18,10 @@
 import java.sql.Connection;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -64,6 +66,12 @@ public Executor(String userName, String password) {
         this.logger = Logger.getLogger(Executor.class.getName()+"2");
     }
 
+    /** Visible for execution tests that provide an in-memory JDBC boundary. */
+    Executor(ConnectionManager connectionManager) {
+        this.connectionManager = connectionManager;
+        this.logger = Logger.getLogger(Executor.class.getName());
+    }
+
     /**
      * run test file function
      */
@@ -106,6 +114,7 @@ public void run(TestScript script) {
         long start = System.currentTimeMillis();
         boolean[] transformed = new boolean[1];
         long waitExpectDeadline = 0;
+        Map<String, String> capturedValues = new HashMap<>();
 
         // for (SqlCommand command : commands) {
         for (int i = 0; i < commands.size(); i++) {
@@ -147,7 +156,7 @@ public void run(TestScript script) {
                 continue;
             }
 
-            String sqlCmd = transformFlushSql(cmd, transformed);
+            String sqlCmd = transformFlushSql(replaceCapturedValues(cmd, capturedValues), transformed);
             if (transformed[0]) {
                 connection = connectionManager.getConnectionForSys();
             } else {
@@ -223,6 +232,16 @@ public void run(TestScript script) {
                 }
 
                 ResultSet resultSet = statement.getResultSet();
+                if (command.getCaptureName() != null) {
+                    captureScalar(command, resultSet, capturedValues);
+                    StmtResult captureResult = new StmtResult();
+                    command.setActResult(captureResult);
+                    command.getTestResult().setActResult(captureResult.toString());
+                    script.addSuccessCmd(command);
+                    command.getTestResult().setResult(RESULT.RESULT_TYPE_PASS);
+                    statement.close();
+                    continue;
+                }
                 if (resultSet != null) {
                     RSSet rsSet = new RSSet(resultSet, command);
                     StmtResult actResult = new StmtResult(rsSet);
@@ -264,6 +283,14 @@ public void run(TestScript script) {
 
             } catch (SQLException e) {
                 waitExpectDeadline = 0;
+                // @capture controls subsequent SQL expansion.  It must never be
+                // compared as an ordinary expected SQL error, otherwise an
+                // unresolved {{name}} can make a broken script look successful.
+                if (command.getCaptureName() != null) {
+                    failCapture(command, script, e);
+                    closeStatement(statement);
+                    continue;
+                }
                 try {
                     if (connection.isClosed() || !connection.isValid(10)) {
                         logger.error("[" + script.getFileName() + "][row:" + command.getPosition() + "]["
@@ -365,7 +392,8 @@ public boolean genRS(TestScript script) {
         connectionManager.reset();
         Connection connection = connectionManager.getConnection();
 
-        Statement statement;
+        Statement statement = null;
+        boolean captureFailed = false;
         // check whether the result file exists
         File rsf = new File(script.getFileName().replaceAll("\\.[A-Za-z]+", COMMON.R_FILE_SUFFIX));
 
@@ -383,6 +411,7 @@ public boolean genRS(TestScript script) {
 
         try (BufferedWriter rs_writer = new BufferedWriter(new FileWriter(rsf.getPath()))) {
             ArrayList<SqlCommand> commands = script.getCommands();
+            Map<String, String> capturedValues = new HashMap<>();
             for (int j = 0; j < commands.size(); j++) {
                 SqlCommand command = null;
 
@@ -420,7 +449,8 @@ public boolean genRS(TestScript script) {
                     connection = getConnection(command);
                     statement = connection.createStatement();
 
-                    String sqlCmd = command.getCommand().replaceAll("\\$resources", COMMON.RESOURCE_PATH);
+                    String sqlCmd = replaceCapturedValues(command.getCommand(), capturedValues)
+                            .replaceAll("\\$resources", COMMON.RESOURCE_PATH);
                     if (command.isNeedWait()) {
                         execWaitOperation(command);
                         // In genRS mode, wait for the condition to be met before executing
@@ -462,6 +492,13 @@ public boolean genRS(TestScript script) {
                         }
                     }
                     ResultSet resultSet = statement.getResultSet();
+                    if (command.getCaptureName() != null) {
+                        captureScalar(command, resultSet, capturedValues);
+                        rs_writer.write(command.getCommand().trim());
+                        if (j < commands.size() - 1) rs_writer.newLine();
+                        statement.close();
+                        continue;
+                    }
                     if (resultSet != null) {
                         RSSet rsSet = new RSSet(resultSet, command);
                         StmtResult actResult = new StmtResult(rsSet);
@@ -486,6 +523,12 @@ public boolean genRS(TestScript script) {
                     }
                     statement.close();
                 } catch (SQLException e) {
+                    if (command != null && command.getCaptureName() != null) {
+                        logger.error("@capture failed; refusing to generate a result file entry", e);
+                        closeStatement(statement);
+                        captureFailed = true;
+                        break;
+                    }
                     rs_writer.write(command.getCommand().trim());
                     rs_writer.newLine();
                     writeRegexPatterns(rs_writer, command);
@@ -501,21 +544,68 @@ public boolean genRS(TestScript script) {
                     throw new RuntimeException(e);
                 }
             }
-            rs_writer.newLine();
-            rs_writer.flush();
-            // drop the test db
-            dropTestDB(connection, script.getUseDB());
+            if (!captureFailed) {
+                rs_writer.newLine();
+                rs_writer.flush();
+            }
         } catch (IOException e) {
             e.printStackTrace();
             return false;
+        } finally {
+            closeStatement(statement);
+            dropTestDB(connection, script.getUseDB());
         }
-        return true;
+        return !captureFailed;
     }
 
     public Connection getConnection(SqlCommand command) {
         return connectionManager.getConnection(command.getConn_id(), getConnUser(command), getConnPswd(command));
     }
 
+    private static String replaceCapturedValues(String sql, Map<String, String> capturedValues) {
+        String result = sql;
+        for (Map.Entry<String, String> entry : capturedValues.entrySet()) {
+            result = result.replace("{{" + entry.getKey() + "}}", entry.getValue());
+        }
+        return result;
+    }
+
+    private static void captureScalar(SqlCommand command, ResultSet resultSet,
+                                      Map<String, String> capturedValues) throws SQLException {
+        if (resultSet == null) {
+            throw new SQLException("@capture requires a query returning one non-null scalar value");
+        }
+        try {
+            if (!resultSet.next() || resultSet.getObject(1) == null) {
+                throw new SQLException("@capture requires a query returning one non-null scalar value");
+            }
+            capturedValues.put(command.getCaptureName(), resultSet.getString(1));
+        } finally {
+            resultSet.close();
+        }
+    }
+
+    private void failCapture(SqlCommand command, TestScript script, SQLException error) {
+        StmtResult result = new StmtResult();
+        result.setType(RESULT.STMT_RESULT_TYPE_ERROR);
+        result.setErrorMessage(error.getMessage());
+        command.setActResult(result);
+        command.getTestResult().setActResult(result.toString());
+        command.getTestResult().setResult(RESULT.RESULT_TYPE_FAILED);
+        script.addFailedCmd(command);
+        logger.error("[" + script.getFileName() + "][row:" + command.getPosition()
+                + "] @capture failed: " + error.getMessage());
+    }
+
+    private static void closeStatement(Statement statement) {
+        if (statement == null) return;
+        try {
+            statement.close();
+        } catch (SQLException ignored) {
+            // The capture failure is the actionable error.
+        }
+    }
+
     private String getConnUser(SqlCommand command) {
         return command.getConn_user() == null ? connectionManager.getDefaultUserName() : command.getConn_user();
     }

src/main/java/io/mo/util/ScriptParser.java

@@ -122,6 +122,12 @@ private boolean processCommentLine(String trimmedLine, String path, SqlCommand c
             parseRegexFlag(trimmedLine.substring(COMMON.REGEX_FLAG.length()), command);
         } else if (trimmedLine.startsWith(COMMON.WAIT_EXPECT_FLAG)) {
             parseWaitExpectFlag(trimmedLine, command);
+        } else if (trimmedLine.startsWith(COMMON.CAPTURE_FLAG)) {
+            String name = trimmedLine.substring(COMMON.CAPTURE_FLAG.length()).trim();
+            if (!name.matches("[A-Za-z_][A-Za-z0-9_]*")) {
+                throw new RuntimeException("Invalid @capture name: " + name);
+            }
+            command.setCaptureName(name);
         }
 
         return false;
@@ -573,4 +579,3 @@ public static void main(String[] args){
 
     }
 }
-