fix: Use run_async during update_test_file

Signed-off-by: Rohan Krishnaswamy <ro@rohank.me>

Author: rkrishn7

sqllogictest/src/parser.rs

@@ -206,6 +206,16 @@ impl<T: ColumnType> Record<T> {
     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

sqllogictest/src/runner.rs

@@ -959,18 +959,11 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
         &mut self,
         record: Record<D::ColumnType>,
     ) -> Result<RecordOutput<D::ColumnType>, 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,16 @@ impl<D: AsyncDB, M: MakeConnection<Conn = D>> Runner<D, M> {
                         );
                         continue;
                     }
+
+                    // If a retry configuration is specified for the record, sleep for the maximum
+                    // retry time.
+                    if let Some(config) = record.retry_config() {
+                        let max_duration = config.backoff * config.attempts.try_into()?;
+                        D::sleep(max_duration).await;
+                    }
+
                     let record_output = self.apply_record(record.clone()).await;
+
                     let record = update_record_with_output(
                         &record,
                         &record_output,