diff --git a/sqllogictest/src/parser.rs b/sqllogictest/src/parser.rs index 475380cc..c7c4211b 100644 --- a/sqllogictest/src/parser.rs +++ b/sqllogictest/src/parser.rs @@ -206,6 +206,16 @@ impl Record { pub fn unparse(&self, w: &mut impl std::io::Write) -> std::io::Result<()> { write!(w, "{self}") } + + /// Returns the retry configuration for this record, if one exists. + pub fn retry_config(&self) -> Option<&RetryConfig> { + match &self { + Record::Statement { retry, .. } + | Record::Query { retry, .. } + | Record::System { retry, .. } => retry.as_ref(), + _ => None, + } + } } /// As is the standard for Display, does not print any trailing diff --git a/sqllogictest/src/runner.rs b/sqllogictest/src/runner.rs index 60b8cc69..678e717a 100644 --- a/sqllogictest/src/runner.rs +++ b/sqllogictest/src/runner.rs @@ -959,18 +959,11 @@ impl> Runner { &mut self, record: Record, ) -> Result, TestError> { - let retry = match &record { - Record::Statement { retry, .. } => retry.clone(), - Record::Query { retry, .. } => retry.clone(), - Record::System { retry, .. } => retry.clone(), - _ => None, - }; - if retry.is_none() { + // The parser ensures that `retry.attempts` must be > 0. + let Some(retry) = record.retry_config() else { return self.run_async_no_retry(record).await; - } + }; - // Retry for `retry.attempts` times. The parser ensures that `retry.attempts` must > 0. - let retry = retry.unwrap(); let mut last_error = None; for _ in 0..retry.attempts { let result = self.run_async_no_retry(record.clone()).await; @@ -1530,7 +1523,18 @@ impl> Runner { ); continue; } - let record_output = self.apply_record(record.clone()).await; + + // If a retry configuration is specified, it is safe run the record with retries + let record_output = if record.retry_config().is_some() { + // Exhaust all possible retries first before applying the record + match self.run_async(record.clone()).await { + Ok(record_output) => record_output, + Err(_) => self.apply_record(record.clone()).await, + } + } else { + self.apply_record(record.clone()).await + }; + let record = update_record_with_output( &record, &record_output,