impl

risinglightdb · Jan 3, 2025 · 35c0467 · 35c0467
1 parent c3b9d45
commit 35c0467
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 15 deletions.
diff --git a/sqllogictest/src/parser.rs b/sqllogictest/src/parser.rs
@@ -755,12 +755,16 @@ fn parse_inner<T: ColumnType>(loc: &Location, script: &str) -> Result<Vec<Record
             ["statement", res @ ..] => {
                 let (mut expected, res) = match res {
                     ["ok", retry @ ..] => (StatementExpect::Ok, retry),
-                    ["error", err_tokens @ ..] => {
-                        // NOTE: `statement error` can't be used with `retry` now because all the
-                        // tokens after `error` are treated as error message.
-                        let error = ExpectedError::parse_inline_tokens(err_tokens)
-                            .map_err(|e| e.at(loc.clone()))?;
-                        (StatementExpect::Error(error), &[][..])
+                    ["error", res @ ..] => {
+                        if res.len() == 4 && res[0] == "retry" && res[2] == "backoff" {
+                            // `statement error retry <num> backoff <duration>`
+                            // To keep syntax simple, let's assume the error message must be multiline.
+                            (StatementExpect::Error(ExpectedError::Empty), res)
+                        } else {
+                            let error = ExpectedError::parse_inline_tokens(res)
+                                .map_err(|e| e.at(loc.clone()))?;
+                            (StatementExpect::Error(error), &[][..])
+                        }
                     }
                     ["count", count_str, retry @ ..] => {
                         let count = count_str.parse::<u64>().map_err(|_| {
@@ -799,12 +803,16 @@ fn parse_inner<T: ColumnType>(loc: &Location, script: &str) -> Result<Vec<Record
             }
             ["query", res @ ..] => {
                 let (mut expected, res) = match res {
-                    ["error", err_tokens @ ..] => {
-                        // NOTE: `query error` can't be used with `retry` now because all the tokens
-                        // after `error` are treated as error message.
-                        let error = ExpectedError::parse_inline_tokens(err_tokens)
-                            .map_err(|e| e.at(loc.clone()))?;
-                        (QueryExpect::Error(error), &[][..])
+                    ["error", res @ ..] => {
+                        if res.len() == 4 && res[0] == "retry" && res[2] == "backoff" {
+                            // `query error retry <num> backoff <duration>`
+                            // To keep syntax simple, let's assume the error message must be multiline.
+                            (QueryExpect::Error(ExpectedError::Empty), res)
+                        } else {
+                            let error = ExpectedError::parse_inline_tokens(res)
+                                .map_err(|e| e.at(loc.clone()))?;
+                            (QueryExpect::Error(error), &[][..])
+                        }
                     }
                     [type_str, res @ ..] => {
                         let types = type_str

diff --git a/tests/retry/query_retry.slt b/tests/retry/query_retry.slt
@@ -10,7 +10,17 @@ SELECT id FROM test ORDER BY random();
 2
 3
 
-query I retry 1 backoff 500ms
+query I retry 2 backoff 500ms
 SELECT id FROM test;
 ----
-1 
+1 
+
+query error retry 2 backoff 500ms
+SELECT id FROM test;
+----
+table not found
+
+query error
+SELECT id FROM test;
+----
+table not found
diff --git a/tests/retry/statement_retry.slt b/tests/retry/statement_retry.slt
@@ -2,4 +2,14 @@ statement ok retry 3 backoff 5s
 INSERT INTO test VALUES (1);
 
 statement count 5 retry 2 backoff 1s
-UPDATE test SET value = value + 1; 
+UPDATE test SET value = value + 1; 
+
+statement error retry 2 backoff 500ms
+UPDATE test SET value = value + 1; 
+----
+table not found
+
+statement error
+UPDATE test SET value = value + 1; 
+----
+table not found