From 436884afdab83d6ba62e6e662ad05ff9a4d72f03 Mon Sep 17 00:00:00 2001
From: Noel Kwan <47273164+kwannoel@users.noreply.github.com>
Date: Fri, 5 Apr 2024 22:52:56 +0800
Subject: [PATCH] chore(sqlsmith): skip timeout queries (#16160)

---
 src/tests/sqlsmith/src/test_runners/diff.rs  |  4 ++--
 src/tests/sqlsmith/src/test_runners/utils.rs | 18 +++++++++++++-----
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/tests/sqlsmith/src/test_runners/diff.rs b/src/tests/sqlsmith/src/test_runners/diff.rs
index f05c9f3521d79..e8ecd0ee89a6b 100644
--- a/src/tests/sqlsmith/src/test_runners/diff.rs
+++ b/src/tests/sqlsmith/src/test_runners/diff.rs
@@ -100,13 +100,13 @@ async fn diff_stream_and_batch_with_sqls(
 
     let select = format!("SELECT * FROM {}", &mview_name);
     tracing::info!("[RUN SELECT * FROM MVIEW id={}]: {}", i, select);
-    let (skip_count, stream_result) = run_query_inner(12, client, &select).await?;
+    let (skip_count, stream_result) = run_query_inner(12, client, &select, true).await?;
     if skip_count > 0 {
         bail!("SQL should not fail: {:?}", select)
     }
 
     tracing::info!("[RUN - BATCH QUERY id={}]: {}", i, &batch);
-    let (skip_count, batch_result) = run_query_inner(12, client, batch).await?;
+    let (skip_count, batch_result) = run_query_inner(12, client, batch, true).await?;
     if skip_count > 0 {
         tracing::info!(
             "[DIFF - DROP MVIEW id={}]: {}",
diff --git a/src/tests/sqlsmith/src/test_runners/utils.rs b/src/tests/sqlsmith/src/test_runners/utils.rs
index 98f29df490446..ab8d57c9b4695 100644
--- a/src/tests/sqlsmith/src/test_runners/utils.rs
+++ b/src/tests/sqlsmith/src/test_runners/utils.rs
@@ -283,9 +283,10 @@ pub(super) fn validate_response(
 }
 
 pub(super) async fn run_query(timeout_duration: u64, client: &Client, query: &str) -> Result<i64> {
-    let (skipped_count, _) = run_query_inner(timeout_duration, client, query).await?;
+    let (skipped_count, _) = run_query_inner(timeout_duration, client, query, true).await?;
     Ok(skipped_count)
 }
+
 /// Run query, handle permissible errors
 /// For recovery error, just do bounded retry.
 /// For other errors, validate them accordingly, skipping if they are permitted.
@@ -296,15 +297,22 @@ pub(super) async fn run_query_inner(
     timeout_duration: u64,
     client: &Client,
     query: &str,
+    skip_timeout: bool,
 ) -> Result<(i64, Vec<SimpleQueryMessage>)> {
     let query_task = client.simple_query(query);
     let result = timeout(Duration::from_secs(timeout_duration), query_task).await;
     let response = match result {
         Ok(r) => r,
-        Err(_) => bail!(
-            "[UNEXPECTED ERROR] Query timeout after {timeout_duration}s:\n{:?}",
-            query
-        ),
+        Err(_) => {
+            if skip_timeout {
+                return Ok((1, vec![]));
+            } else {
+                bail!(
+                    "[UNEXPECTED ERROR] Query timeout after {timeout_duration}s:\n{:?}",
+                    query
+                )
+            }
+        }
     };
     if let Err(e) = &response
         && let Some(e) = e.as_db_error()