From 35c046790cd08aa3e1524f6e209efb53c6ad3053 Mon Sep 17 00:00:00 2001
From: Eric Fu <fuyufjh@gmail.com>
Date: Fri, 3 Jan 2025 16:28:14 +0800
Subject: [PATCH] impl

---
 sqllogictest/src/parser.rs      | 32 ++++++++++++++++++++------------
 tests/retry/query_retry.slt     | 14 ++++++++++++--
 tests/retry/statement_retry.slt | 12 +++++++++++-
 3 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/sqllogictest/src/parser.rs b/sqllogictest/src/parser.rs
index 9e0b2ce..b069469 100644
--- 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
index 0c030ab..209bf40 100644
--- 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 
\ No newline at end of file
+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
index 6559ce6..110f5d0 100644
--- 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; 
\ No newline at end of file
+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