feat: support retry for [statement|query] error (#244)

* impl Signed-off-by: Eric Fu <[email protected]> * readme Signed-off-by: Eric Fu <[email protected]> * minor fixes Signed-off-by: Eric Fu <[email protected]> --------- Signed-off-by: Eric Fu <[email protected]>
risinglightdb · Jan 3, 2025 · 99e8301 · 99e8301
1 parent c3b9d45
commit 99e8301
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 15 deletions.
diff --git a/README.md b/README.md
@@ -140,10 +140,23 @@ SELECT id FROM test;
 ----
 1
 
+query error retry 3 backoff 5s
+SELECT id FROM test;
+----
+database error: table not found
+
+
 statement ok retry 3 backoff 5s
 UPDATE test SET id = 1;
+
+statement error
+UPDATE test SET value = value + 1; 
+----
+database error: table not found
 ```
 
+Due to the limitation of syntax, the retry clause can't be used along with the single-line regex error message extension.
+
 ### Extension: Environment variable substitution in query and statement
 
 It needs to be enabled by adding `control substitution on` to the test file.

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,19 @@ 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,16 @@ 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
+