diff --git a/README.md b/README.md index 2686c9b08..38a438265 100644 --- a/README.md +++ b/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.}
-- @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.

**Example:**
`-- @wait_expect(1, 10)`
`SELECT * FROM orders WHERE status = 'completed';`

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.

**Modes:**
- **run mode**: Polls until result matches or timeout
- **genrs mode**: Sleeps for timeout seconds, then executes once

**Features:**
- Supports query results (SELECT)
- Supports error matching (waiting for errors to appear/disappear)
- Detailed logging of retry attempts and timing | | -- @session:id=2&user=root&password=111
-- @session | The mo-tester will create a new connetion to execute sql statements between those two tags.
Default value of id is 1, max is 10.
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.
-- @sortkey:0,1: means sort keys are first column and second colum. | diff --git a/src/main/java/io/mo/cases/SqlCommand.java b/src/main/java/io/mo/cases/SqlCommand.java index 569db42cc..147d950e9 100755 --- a/src/main/java/io/mo/cases/SqlCommand.java +++ b/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(); diff --git a/src/main/java/io/mo/constant/COMMON.java b/src/main/java/io/mo/constant/COMMON.java index 30d24967c..8f7ec2a92 100755 --- a/src/main/java/io/mo/constant/COMMON.java +++ b/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; -} \ No newline at end of file +} diff --git a/src/main/java/io/mo/db/Executor.java b/src/main/java/io/mo/db/Executor.java index ee5719b73..11321a7c8 100755 --- a/src/main/java/io/mo/db/Executor.java +++ b/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 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 commands = script.getCommands(); + Map 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 capturedValues) { + String result = sql; + for (Map.Entry entry : capturedValues.entrySet()) { + result = result.replace("{{" + entry.getKey() + "}}", entry.getValue()); + } + return result; + } + + private static void captureScalar(SqlCommand command, ResultSet resultSet, + Map 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(); } diff --git a/src/main/java/io/mo/util/ScriptParser.java b/src/main/java/io/mo/util/ScriptParser.java index 5f64933cc..aee0f82d7 100755 --- a/src/main/java/io/mo/util/ScriptParser.java +++ b/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){ } } - diff --git a/src/test/java/io/mo/db/ExecutorCaptureTest.java b/src/test/java/io/mo/db/ExecutorCaptureTest.java new file mode 100644 index 000000000..e0bfecd98 --- /dev/null +++ b/src/test/java/io/mo/db/ExecutorCaptureTest.java @@ -0,0 +1,166 @@ +package io.mo.db; + +import io.mo.cases.SqlCommand; +import io.mo.cases.TestScript; +import org.junit.Test; + +import java.io.File; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +/** Exercises capture failures through Executor.genRS, not just parser state. */ +public class ExecutorCaptureTest { + + @Test + public void captureQueryErrorAbortsResultGeneration() throws Exception { + assertCaptureFailure(FailureMode.QUERY_ERROR); + } + + @Test + public void captureEmptyResultAbortsResultGeneration() throws Exception { + assertCaptureFailure(FailureMode.EMPTY_RESULT); + } + + @Test + public void captureNullResultAbortsResultGeneration() throws Exception { + assertCaptureFailure(FailureMode.NULL_VALUE); + } + + private void assertCaptureFailure(FailureMode mode) throws Exception { + File sql = File.createTempFile("capture-failure", ".sql"); + File result = new File(sql.getPath().replaceAll("\\.[A-Za-z]+", ".result")); + try { + TestScript script = new TestScript(); + script.setFileName(sql.getPath()); + SqlCommand command = new SqlCommand(); + command.append("select capture_value();\n"); + command.setCaptureName("value"); + script.addCommand(command); + + Tracker tracker = new Tracker(); + boolean generated = new Executor(new StubConnectionManager(mode, tracker)).genRS(script); + + assertFalse("a control-directive failure must not become expected output", generated); + assertTrue("capture failure must not be written to the result file", + !result.exists() || result.length() == 0); + assertTrue("capture statement must be closed", tracker.captureStatementClosed); + assertTrue("test database must be dropped", tracker.testDatabaseDropped); + if (mode != FailureMode.QUERY_ERROR) { + assertTrue("capture result set must be closed", tracker.captureResultSetClosed); + } + } finally { + sql.delete(); + result.delete(); + } + } + + private enum FailureMode { QUERY_ERROR, EMPTY_RESULT, NULL_VALUE } + + private static final class Tracker { + boolean captureStatementClosed; + boolean captureResultSetClosed; + boolean testDatabaseDropped; + } + + private static final class StubConnectionManager extends ConnectionManager { + private final Connection connection; + + StubConnectionManager(FailureMode mode, Tracker tracker) { + connection = proxy(Connection.class, new ConnectionHandler(mode, tracker)); + } + + @Override public Connection getConnection() { return connection; } + @Override public Connection getConnection(int index) { return connection; } + @Override public Connection getConnection(int index, String user, String password) { return connection; } + @Override public void reset() { } + } + + private static final class ConnectionHandler implements InvocationHandler { + private final FailureMode mode; + private final Tracker tracker; + ConnectionHandler(FailureMode mode, Tracker tracker) { + this.mode = mode; + this.tracker = tracker; + } + + @Override public Object invoke(Object proxy, java.lang.reflect.Method method, Object[] args) { + String name = method.getName(); + if ("createStatement".equals(name)) return proxy(Statement.class, new StatementHandler(mode, tracker)); + if ("isClosed".equals(name)) return false; + if ("isValid".equals(name)) return true; + return defaultValue(method.getReturnType()); + } + } + + private static final class StatementHandler implements InvocationHandler { + private final FailureMode mode; + private final Tracker tracker; + private boolean captureStatement; + StatementHandler(FailureMode mode, Tracker tracker) { + this.mode = mode; + this.tracker = tracker; + } + + @Override public Object invoke(Object proxy, java.lang.reflect.Method method, Object[] args) throws Throwable { + String name = method.getName(); + if ("execute".equals(name)) { + String sql = (String) args[0]; + if (sql.contains("capture_value")) { + captureStatement = true; + if (mode == FailureMode.QUERY_ERROR) throw new SQLException("capture query failed"); + return true; + } + if (sql.startsWith("drop database")) tracker.testDatabaseDropped = true; + return false; + } + if ("getResultSet".equals(name)) { + return proxy(ResultSet.class, new ResultSetHandler(mode, tracker)); + } + if ("close".equals(name) && captureStatement) tracker.captureStatementClosed = true; + return defaultValue(method.getReturnType()); + } + } + + private static final class ResultSetHandler implements InvocationHandler { + private final FailureMode mode; + private final Tracker tracker; + private boolean advanced; + ResultSetHandler(FailureMode mode, Tracker tracker) { + this.mode = mode; + this.tracker = tracker; + } + + @Override public Object invoke(Object proxy, java.lang.reflect.Method method, Object[] args) { + if ("next".equals(method.getName())) { + if (mode == FailureMode.EMPTY_RESULT || advanced) return false; + advanced = true; + return true; + } + if ("getObject".equals(method.getName())) return null; + if ("close".equals(method.getName())) tracker.captureResultSetClosed = true; + return defaultValue(method.getReturnType()); + } + } + + @SuppressWarnings("unchecked") + private static T proxy(Class type, InvocationHandler handler) { + return (T) Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, handler); + } + + private static Object defaultValue(Class type) { + if (!type.isPrimitive()) return null; + if (type == Boolean.TYPE) return false; + if (type == Character.TYPE) return '\0'; + if (type == Byte.TYPE || type == Short.TYPE || type == Integer.TYPE || type == Long.TYPE) return 0; + if (type == Float.TYPE) return 0F; + if (type == Double.TYPE) return 0D; + return null; + } +} diff --git a/src/test/java/io/mo/util/ScriptParserTest.java b/src/test/java/io/mo/util/ScriptParserTest.java index cbde0cb79..2b9a78696 100644 --- a/src/test/java/io/mo/util/ScriptParserTest.java +++ b/src/test/java/io/mo/util/ScriptParserTest.java @@ -193,6 +193,20 @@ public void testSystemCommandFlag() { // 系统命令标记应该被解析 assertTrue("Should have parsed system commands", foundSystemCommand); } + + @Test + public void testCaptureFlag() throws Exception { + File script = File.createTempFile("capture-flag", ".sql"); + script.deleteOnExit(); + try (PrintWriter writer = new PrintWriter(new FileWriter(script))) { + writer.println("-- @capture:restore_ts"); + writer.println("select '2026-08-25 00:00:00';"); + } + + TestScript testScript = new ScriptParser().parseScript(script.getAbsolutePath()); + assertEquals(1, testScript.getCommands().size()); + assertEquals("restore_ts", testScript.getCommands().get(0).getCaptureName()); + } /** * 测试 Session 连接标记 @@ -877,4 +891,3 @@ public void testMixedMetacmpFlags() throws Exception { Boolean.TRUE, testScript.getCommands().get(5).getCompareMeta()); } } -